Update utils/translator.py

utils/translator.py

@@ -2,34 +2,26 @@
 
 from transformers import AutoTokenizer, AutoModelForSeq2SeqLM
 import torch
+from docx import Document
 
-# 🧠 Load formal Brazilian Portuguese legal translator
-model_name = "unicamp-dl/translation-en-pt-t5"
-tokenizer = AutoTokenizer.from_pretrained(model_name)
-model = AutoModelForSeq2SeqLM.from_pretrained(model_name)
+# ========== Load EN → PT model ==========
+en_pt_model_name = "unicamp-dl/translation-en-pt-t5"
+tokenizer_en_pt = AutoTokenizer.from_pretrained(en_pt_model_name)
+model_en_pt = AutoModelForSeq2SeqLM.from_pretrained(en_pt_model_name)
 
-def translate_text(text):
-    if not text.strip():
-        return "No input provided."
-
-    text = text.replace("\n", " ").strip()
-    chunks = [text[i:i + 512] for i in range(0, len(text), 512)]
-    translated = []
-
-    for chunk in chunks:
-        inputs = tokenizer(chunk, return_tensors="pt", truncation=True, padding=True)
-        with torch.no_grad():
-            outputs = model.generate(**inputs, max_length=512, num_beams=4)
-        translated.append(tokenizer.decode(outputs[0], skip_special_tokens=True))
+# ========== Load PT → EN model ==========
+pt_en_model_name = "unicamp-dl/translation-pt-en-t5"
+tokenizer_pt_en = AutoTokenizer.from_pretrained(pt_en_model_name)
+model_pt_en = AutoModelForSeq2SeqLM.from_pretrained(pt_en_model_name)
 
-    return " ".join(translated)
+# ========== Text Cleaning & Chunking ==========
 
 def clean_text(text):
     return text.replace("\n", " ").replace("  ", " ").strip()
 
 def chunk_text(text, max_chunk_chars=500):
     """
-    🔪 Chunk long text into segments that fit model constraints.
+    🔪 Break long input into token-safe chunks.
     """
     words = text.split()
     chunks = []
@@ -46,9 +38,32 @@ def chunk_text(text, max_chunk_chars=500):
     
     return chunks
 
+# ========== Translation Functions ==========
+
 def translate_to_portuguese(text):
     """
-    🌐 Clean, chunk, and translate English text into Portuguese.
+    🇺🇸 ➡️ 🇧🇷 Translate English to Portuguese.
+    """
+    if not text.strip():
+        return "No input provided."
+
+    text = clean_text(text)
+    chunks = chunk_text(text)
+
+    translated_chunks = []
+    for chunk in chunks:
+        inputs = tokenizer_en_pt(chunk, return_tensors="pt", truncation=True, padding=True)
+        with torch.no_grad():
+            outputs = model_en_pt.generate(**inputs, max_length=512, num_beams=4)
+        translated = tokenizer_en_pt.decode(outputs[0], skip_special_tokens=True)
+        translated_chunks.append(translated)
+
+    return " ".join(translated_chunks)
+
+
+def translate_to_english(text):
+    """
+    🇧🇷 ➡️ 🇺🇸 Translate Portuguese to English.
     """
     if not text.strip():
         return "No input provided."
@@ -58,7 +73,37 @@ def translate_to_portuguese(text):
 
     translated_chunks = []
     for chunk in chunks:
-        result = translator(chunk, max_length=512)
-        translated_chunks.append(result[0]["translation_text"])
+        inputs = tokenizer_pt_en(chunk, return_tensors="pt", truncation=True, padding=True)
+        with torch.no_grad():
+            outputs = model_pt_en.generate(**inputs, max_length=512, num_beams=4)
+        translated = tokenizer_pt_en.decode(outputs[0], skip_special_tokens=True)
+        translated_chunks.append(translated)
 
     return " ".join(translated_chunks)
+
+# ========== Bilingual Layout ==========
+
+def bilingual_clauses(english_text):
+    """
+    📄 Generate side-by-side bilingual clauses.
+    """
+    clauses_en = chunk_text(clean_text(english_text), max_chunk_chars=300)
+    clauses_pt = [translate_to_portuguese(c) for c in clauses_en]
+
+    bilingual = []
+    for en, pt in zip(clauses_en, clauses_pt):
+        bilingual.append(f"📘 EN: {en}\n📗 PT: {pt}\n" + "-" * 60)
+    return "\n\n".join(bilingual)
+
+# ========== Export to DOCX ==========
+
+def export_to_word(text, filename="translated_contract.docx"):
+    """
+    📝 Export text block to Word document.
+    """
+    doc = Document()
+    doc.add_heading("Legal Translation Output", level=1)
+    for para in text.split("\n\n"):
+        doc.add_paragraph(para)
+    doc.save(filename)
+    return filename