File size: 2,711 Bytes
5a7ab71
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
"""
Do optional cleaning (e.g., remove some languages).

Usage:
python3 -m fastchat.data.optional_clean --in input.json --out output.json --keep-lang en
python3 -m fastchat.data.optional_clean --in input.json --out output.json --skip-lang en

Requirement:
pip3 install polyglot icu pyicu pycld2 morfessor
"""
import argparse
import json
import re

import polyglot
from polyglot.detect import Detector
import pycld2
from tqdm import tqdm


def skip(conv, args):
    # Remove certain languages
    if args.keep_lang != "all" or args.skip_lang is not None:
        text = "\n".join([x["value"] for x in conv["conversations"]])
        try:
            lang_code = Detector(text).language.code
        except (pycld2.error, polyglot.detect.base.UnknownLanguage):
            lang_code = "unknown"

        if args.keep_lang != "all" and lang_code != args.keep_lang:
            return True

        if lang_code == args.skip_lang:
            return True

    # Remove repetitive numbers
    if args.reduce_rep:
        for sentence in conv["conversations"]:
            val = sentence["value"]
            sub = re.search(r"(\d)\1{8}", val)
            if sub is not None:
                return True

    return False


if __name__ == "__main__":
    parser = argparse.ArgumentParser()
    parser.add_argument("--in-file", type=str, required=True)
    parser.add_argument("--out-file", type=str)
    parser.add_argument(
        "--keep-lang",
        type=str,
        default="all",
        choices=["all", "en"],
        help="Only keep certain langauges.",
    )
    parser.add_argument("--skip-lang", type=str, help="Skip a specific language.")
    # NOTE: Be careful about reduce_rep which may remove some good data.
    # For example, addresses could have long consecutive 0's
    parser.add_argument("--reduce-rep", action="store_true")
    args = parser.parse_args()

    in_file = args.in_file
    out_file = args.out_file
    keep_lang = args.keep_lang
    skip_lang = args.skip_lang
    reduce_rep = args.reduce_rep
    assert keep_lang == "all" or skip_lang is None

    if out_file is None:
        out_file = "sharegpt_clean"
        if keep_lang != "all":
            out_file += "_" + keep_lang
        if skip_lang is not None:
            out_file += "_skip_" + skip_lang
        if reduce_rep:
            out_file += "_reduce_rep"
        out_file += ".json"

    content = json.load(open(in_file, "r"))
    num_conv = len(content)

    new_content = []
    for conv in tqdm(content):
        if not skip(conv, args):
            new_content.append(conv)

    print(f"return {len(new_content)} out of {len(content)}, start dump ...")
    json.dump(new_content, open(out_file, "w"), indent=2)