File size: 3,648 Bytes
7e9cfd1 | 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 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 | """Rewrite tweets with a chat LLM to get a few more training rows per label,
keeping the same dialect and stance. Point it at your endpoint with
AUG_BASE_URL / AUG_MODEL.
python -m src.augment --in data/track2/train.csv \\
--out data/track2/train_aug.csv --n 2
"""
import argparse
import json
import os
import urllib.request
from concurrent.futures import ThreadPoolExecutor
import pandas as pd
PROMPT = (
"أعد صياغة التغريدة التالية بأسلوب مختلف مع الحفاظ التام على نفس "
"الموقف والرأي تجاه الموضوع، ونفس اللهجة، ونفس المعنى. لا تضف أي "
"تعليق أو شرح، وأخرج التغريدة المعاد صياغتها فقط:\n\n{text}"
)
def call(base_url, model, text, n, timeout=120):
body = json.dumps({
"model": model,
"n": n,
"temperature": 1.0,
"top_p": 0.95,
"max_tokens": 160,
"messages": [{"role": "user", "content": PROMPT.format(text=text)}],
}).encode("utf-8")
req = urllib.request.Request(
base_url.rstrip("/") + "/chat/completions",
data=body,
headers={"Content-Type": "application/json"},
)
with urllib.request.urlopen(req, timeout=timeout) as r:
data = json.load(r)
return [c["message"]["content"].strip() for c in data["choices"]]
def clean(variant, original):
v = variant.strip().strip('"').strip()
if "\n" in v:
v = v.split("\n")[-1].strip()
if len(v) < 5 or v == original.strip():
return None
return v
def main():
ap = argparse.ArgumentParser()
ap.add_argument("--in", dest="inp", required=True)
ap.add_argument("--out", required=True)
ap.add_argument("--n", type=int, default=2)
ap.add_argument("--concurrency", type=int, default=32)
ap.add_argument("--max_rows", type=int, default=0)
ap.add_argument(
"--base_url", default=os.environ.get("AUG_BASE_URL", "")
)
ap.add_argument("--model", default=os.environ.get("AUG_MODEL", ""))
args = ap.parse_args()
if not args.base_url or not args.model:
raise SystemExit("set --base_url/--model or AUG_BASE_URL/AUG_MODEL")
df = pd.read_csv(args.inp, keep_default_na=False, encoding="utf-8-sig")
if args.max_rows:
df = df.head(args.max_rows)
def work(row):
try:
variants = call(args.base_url, args.model, row["text"], args.n)
except Exception:
return []
out = []
for v in variants:
c = clean(v, row["text"])
if c:
out.append({
"text": c,
"target": row["target"],
"stance": row["stance"],
"source": "paraphrase",
})
return out
rows = df.to_dict("records")
generated = []
done = 0
with ThreadPoolExecutor(max_workers=args.concurrency) as ex:
for res in ex.map(work, rows):
generated.extend(res)
done += 1
if done % 200 == 0:
print(f" {done}/{len(rows)} rows, "
f"{len(generated)} paraphrases")
orig = df.copy()
orig["source"] = "original"
keep = ["text", "target", "stance", "source"]
combined = pd.concat(
[orig[keep], pd.DataFrame(generated)[keep]], ignore_index=True
)
combined.to_csv(args.out, index=False, encoding="utf-8-sig")
print(f"[write] {len(orig)} original + {len(generated)} paraphrases "
f"= {len(combined)} rows -> {args.out}")
if __name__ == "__main__":
main()
|