File size: 2,530 Bytes
9d3f668 | 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 | #!/usr/bin/env python3
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()
|