| |
| import sqlite3 |
| from transformers import MarianMTModel, MarianTokenizer, pipeline |
|
|
| print("Loading translation model...") |
| model_name = "Helsinki-NLP/opus-mt-nso-en" |
| tok = MarianTokenizer.from_pretrained(model_name) |
| model = MarianMTModel.from_pretrained(model_name) |
|
|
| print("Loading sentiment model...") |
| classifier = pipeline("text-classification", |
| model="distilbert-base-uncased-finetuned-sst-2-english", |
| device=-1) |
| print("Models ready.\n") |
|
|
| def translate_batch(texts): |
| inputs = tok(texts, return_tensors="pt", padding=True, truncation=True, max_length=128) |
| outputs = model.generate(**inputs, max_length=128) |
| return [tok.decode(o, skip_special_tokens=True) for o in outputs] |
|
|
| conn = sqlite3.connect("/home/sediba/sepedi_datasets/sepedi_data.db") |
| done = set(r[0] for r in conn.execute("SELECT text FROM annotated_data").fetchall()) |
|
|
| rows = conn.execute(""" |
| SELECT id, text, source FROM raw_texts |
| WHERE length(text) BETWEEN 20 AND 250 |
| AND text NOT LIKE '%<%' |
| AND text NOT LIKE '%>%' |
| AND source IN ('jw.org', 'SADILAR-NCHLT-Annotated-Corpus') |
| ORDER BY source DESC, RANDOM() |
| LIMIT 2000 |
| """).fetchall() |
|
|
| candidates = [(i, t, s) for i, t, s in rows if t not in done] |
| print(f"Candidates: {len(candidates)}") |
|
|
| inserted, errors = 0, 0 |
| BATCH = 16 |
|
|
| for i in range(0, len(candidates), BATCH): |
| batch = candidates[i:i+BATCH] |
| texts = [r[1] for r in batch] |
| try: |
| eng = translate_batch(texts) |
| sentiments = classifier(eng, truncation=True, max_length=128) |
| for (_, orig, src), sent in zip(batch, sentiments): |
| label = sent["label"].lower() |
| conn.execute(""" |
| INSERT OR IGNORE INTO annotated_data |
| (text, label, label_name, annotator, confidence, date_annotated) |
| VALUES (?, ?, ?, ?, ?, datetime('now')) |
| """, (orig, 1 if label=="positive" else 0, label, |
| "auto:opus-mt+distilbert", round(sent["score"], 4))) |
| inserted += 1 |
| conn.commit() |
| if i % (BATCH*5) == 0: |
| print(f" {inserted} labeled...") |
| except Exception as e: |
| errors += 1 |
| print(f" batch error: {e}") |
|
|
| total = conn.execute("SELECT COUNT(*) FROM annotated_data").fetchone()[0] |
| dist = conn.execute("SELECT label_name, COUNT(*) FROM annotated_data GROUP BY label_name").fetchall() |
| print(f"\nDone. Inserted={inserted} Errors={errors} Total={total}") |
| for l, c in dist: print(f" {l}: {c}") |
| conn.close() |
|
|
