ok

seq2seq_inference.py

@@ -1,5 +1,5 @@
-import json
 import torch
+import json
 from transformers import AutoTokenizer, AutoModelForSeq2SeqLM
 from peft import PeftModel
 
@@ -18,10 +18,38 @@ SYSTEM_PROMPT = (
 )
 
 
+MODEL_OPTIONS = {
+    "Araberta": {
+        "base": "asmashayea/absa-araberta",
+        "adapter": "asmashayea/absa-araberta"
+    },
+    "mT5": {
+        "base": "google/mt5-base",
+        "adapter": "asmashayea/mt4-absa"
+    },
+    "mBART": {
+        "base": "facebook/mbart-large-50-many-to-many-mmt",
+        "adapter": "asmashayea/mbart-absa"
+    },
+}
+
+cached_models = {}
+
+# ✅ Reusable for both mT5 + mBART
+def load_mt5_bart(model_key):
+    base_id = MODEL_OPTIONS[model_key]["base"]
+    adapter_id = MODEL_OPTIONS[model_key]["adapter"]
 
-def infer_t5_prompt(review_text, tokenizer, peft_model):
-    prompt = SYSTEM_PROMPT + f"\n\nReview: {review_text}"
+    tokenizer = AutoTokenizer.from_pretrained(adapter_id)
+    base_model = AutoModelForSeq2SeqLM.from_pretrained(base_id)
+    peft_model = PeftModel.from_pretrained(base_model, adapter_id)
+    peft_model.eval()
 
+    cached_models[model_key] = (tokenizer, peft_model)
+    return tokenizer, peft_model
+
+def infer_t5_bart(text, tokenizer, peft_model):
+    prompt = SYSTEM_PROMPT + f"\n\nReview: {text}"
     inputs = tokenizer(prompt, return_tensors="pt", padding=True, truncation=True).to(peft_model.device)
 
     with torch.no_grad():
@@ -36,15 +64,11 @@ def infer_t5_prompt(review_text, tokenizer, peft_model):
             eos_token_id=tokenizer.eos_token_id,
         )
 
-    decoded = tokenizer.decode(
-        outputs[0],
-        skip_special_tokens=True,
-        clean_up_tokenization_spaces=False
-    ).strip()
-
+    decoded = tokenizer.decode(outputs[0], skip_special_tokens=True).strip()
     decoded = decoded.replace('<extra_id_0>', '').replace('</s>', '').strip()
 
     try:
         return json.loads(decoded)
     except json.JSONDecodeError:
-        return decoded
+        return {"raw_output": decoded, "error": "Invalid JSON"}
+