b

inference.py

@@ -1,8 +1,13 @@
 import torch
-from transformers import AutoTokenizer, AutoModelForSeq2SeqLM
+from transformers import (
+    AutoTokenizer,
+    AutoModelForSeq2SeqLM,
+    MBartForConditionalGeneration,
+    MBart50TokenizerFast
+)
 from peft import PeftModel
 
-# Define supported models and their adapter IDs
+# Supported models and their adapter IDs
 MODEL_OPTIONS = {
     "mT5": {
         "base": "google/mt5-base",
@@ -12,12 +17,13 @@ MODEL_OPTIONS = {
         "base": "facebook/mbart-large-50-many-to-many-mmt",
         "adapter": "asmashayea/mbart-absa"
     },
+    # You can customize GPT-like entries later
     "GPT3.5": {
-        "base": "bigscience/bloom-560m",  # placeholder
+        "base": "bigscience/bloom-560m",  # Placeholder only
         "adapter": "asmashayea/gpt-absa"
     },
     "GPT4o": {
-        "base": "bigscience/bloom-560m",  # placeholder
+        "base": "bigscience/bloom-560m",  # Placeholder only
         "adapter": "asmashayea/gpt-absa"
     }
 }
@@ -31,8 +37,14 @@ def load_model(model_key):
     base_id = MODEL_OPTIONS[model_key]["base"]
     adapter_id = MODEL_OPTIONS[model_key]["adapter"]
 
-    tokenizer = AutoTokenizer.from_pretrained(adapter_id)
-    base_model = AutoModelForSeq2SeqLM.from_pretrained(base_id)
+    if model_key == "mBART":
+        tokenizer = MBart50TokenizerFast.from_pretrained(base_id)
+        tokenizer.src_lang = "ar_AR"  # Required for input
+        base_model = MBartForConditionalGeneration.from_pretrained(base_id)
+    else:
+        tokenizer = AutoTokenizer.from_pretrained(adapter_id)
+        base_model = AutoModelForSeq2SeqLM.from_pretrained(base_id)
+
     model = PeftModel.from_pretrained(base_model, adapter_id)
     model.eval()
 
@@ -41,22 +53,25 @@ def load_model(model_key):
 
 def predict_absa(text, model_choice):
     tokenizer, model = load_model(model_choice)
-    prompt = f"استخرج الجوانب والآراء والمشاعر من النص التالي:\n{text}"
+    prompt = f"استخرج الجوانب والآراء والمشاعر من النص التالي:\n{text.strip()}"
 
     if model_choice == "mBART":
-        # ✅ Critical fix
-        tokenizer.src_lang = "ar_AR"
-        inputs = tokenizer(prompt, return_tensors="pt", truncation=True)
-        forced_bos_token_id = tokenizer.lang_code_to_id["ar_AR"]
-        outputs = model.generate(
-            **inputs,
-            max_new_tokens=128,
-            forced_bos_token_id=forced_bos_token_id  # Force Arabic generation
-        )
+        tokenizer.tgt_lang = "ar_AR"  # Required for output
+        inputs = tokenizer(prompt, return_tensors="pt", padding="max_length", truncation=True, max_length=512)
+
+        with torch.no_grad():
+            outputs = model.generate(
+                **inputs,
+                max_new_tokens=128,
+                do_sample=False,
+                temperature=0.0,
+                forced_bos_token_id=tokenizer.lang_code_to_id["ar_AR"],
+                pad_token_id=tokenizer.pad_token_id,
+            )
     else:
         inputs = tokenizer(prompt, return_tensors="pt", truncation=True)
-        outputs = model.generate(**inputs, max_new_tokens=128)
+        with torch.no_grad():
+            outputs = model.generate(**inputs, max_new_tokens=128)
 
-    decoded = tokenizer.decode(outputs[0], skip_special_tokens=True)
+    decoded = tokenizer.decode(outputs[0], skip_special_tokens=True).strip()
     return decoded
-