peft

inference.py

@@ -1,13 +1,8 @@
 import torch
-from transformers import (
-    AutoTokenizer,
-    AutoModelForSeq2SeqLM,
-    MBartForConditionalGeneration,
-    MBart50TokenizerFast
-)
+from transformers import AutoTokenizer, AutoModelForSeq2SeqLM
 from peft import PeftModel
-
-# Supported models and their adapter IDs
+from seq2seq_inference import infer_t5_prompt, infer_mBart_prompt
+# Define supported models and their adapter IDs
 MODEL_OPTIONS = {
     "mT5": {
         "base": "google/mt5-base",
@@ -17,13 +12,12 @@ 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 only
+        "base": "bigscience/bloom-560m",  # example, not ideal for ABSA
         "adapter": "asmashayea/gpt-absa"
     },
     "GPT4o": {
-        "base": "bigscience/bloom-560m",  # Placeholder only
+        "base": "bigscience/bloom-560m",  # example, not ideal for ABSA
         "adapter": "asmashayea/gpt-absa"
     }
 }
@@ -37,14 +31,8 @@ def load_model(model_key):
     base_id = MODEL_OPTIONS[model_key]["base"]
     adapter_id = MODEL_OPTIONS[model_key]["adapter"]
 
-    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)
-
+    tokenizer = AutoTokenizer.from_pretrained(adapter_id)
+    base_model = AutoModelForSeq2SeqLM.from_pretrained(base_id)
     model = PeftModel.from_pretrained(base_model, adapter_id)
     model.eval()
 
@@ -52,26 +40,21 @@ def load_model(model_key):
     return tokenizer, model
 
 def predict_absa(text, model_choice):
+
     tokenizer, model = load_model(model_choice)
-    prompt = f"استخرج الجوانب والآراء والمشاعر من النص التالي:\n{text.strip()}"
 
-    if model_choice == "mBART":
-        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)
-        with torch.no_grad():
-            outputs = model.generate(**inputs, max_new_tokens=128)
+    if model_choice == 'mT5':
+        decoded = infer_t5_prompt(text, tokenizer, model)
+
+    elif model_choice == 'mBART':
+        decoded = infer_mBart_prompt(text, tokenizer, model)
+    
+    # prompt = f"استخرج الجوانب والآراء والمشاعر من النص التالي:\n{text}"
+    # inputs = tokenizer(prompt, return_tensors="pt", truncation=True)
+
+    # with torch.no_grad():
+    #     outputs = model.generate(**inputs, max_new_tokens=128)
 
-    decoded = tokenizer.decode(outputs[0], skip_special_tokens=True).strip()
+    # decoded = tokenizer.decode(outputs[0], skip_special_tokens=True)
     return decoded

seq2seq_inference.py

@@ -0,0 +1,78 @@
+import json
+import torch
+from transformers import AutoTokenizer, AutoModelForSeq2SeqLM
+from peft import PeftModel
+
+# Updated Prediction Function for mBART
+SYSTEM_PROMPT = (
+    "You are an advanced AI model specialized in extracting aspects and determining their sentiment polarity from customer reviews.\n\n"
+    "Instructions:\n"
+    "1. Extract only the aspects (nouns) mentioned in the review.\n"
+    "2. Assign a sentiment to each aspect: \"positive\", \"negative\", or \"neutral\".\n"
+    "3. Return aspects in the same language as they appear.\n"
+    "4. An aspect must be a noun that refers to a specific item or service the user described.\n"
+    "5. Ignore adjectives, general ideas, and vague topics.\n"
+    "6. Do NOT translate, explain, or add extra text.\n"
+    "7. The output must be just a valid JSON list with 'aspect' and 'sentiment'. Start with `[` and stop at `]`.\n"
+    "8. Do NOT output the instructions, review, or any text — only one output JSON list.\n"
+    "9. Just one output and one review."
+)
+
+def infer_mBart_prompt(review_text, tokenizer, model):
+    # Set target language for mBART
+    tokenizer.tgt_lang = "ar_AR"  # Change as needed ("en_XX" for English)
+
+    prompt = f"{SYSTEM_PROMPT}\nReview: {review_text}"
+
+    inputs = tokenizer(prompt, return_tensors="pt", padding="max_length", truncation=True, max_length=512).to(model.device)
+
+    with torch.no_grad():
+        outputs = model.generate(
+            **inputs,
+            max_new_tokens=128,
+            do_sample=False,
+            temperature=0.0,
+            forced_bos_token_id=tokenizer.convert_tokens_to_ids(tokenizer.tgt_lang),  # safer
+            pad_token_id=tokenizer.pad_token_id
+        )
+
+    decoded = tokenizer.decode(outputs[0], skip_special_tokens=True).replace("ar_AR ", "").replace("en_XX ", "").strip()
+    return decoded
+
+
+
+
+def infer_t5_prompt(review_text, tokenizer, model):
+    prompt = (
+        SYSTEM_PROMPT + f"\n\nReview: {review_text}"
+    )
+    inputs = tokenizer(prompt, return_tensors="pt", padding=True, truncation=True).to(peft_model.device)
+
+    with torch.no_grad():
+        outputs = model.generate(
+            **inputs,
+            max_new_tokens=256,
+            num_beams=4,
+            do_sample=False,
+            temperature=0.0,
+            early_stopping=True,
+            pad_token_id=tokenizer.pad_token_id,
+            eos_token_id=tokenizer.eos_token_id,
+        )
+
+    decoded = tokenizer.decode(
+        outputs[0], 
+        skip_special_tokens=True, 
+        clean_up_tokenization_spaces=False
+    ).strip()
+
+    # Optional: remove T5 special tokens like <extra_id_0>
+    decoded = decoded.replace('<extra_id_0>', '').replace('</s>', '').strip()
+
+    try:
+        json_output = json.loads(decoded)
+    except json.JSONDecodeError as e:
+        print(f"⚠️ JSON decode error: {e}. Returning raw output.")
+        json_output = decoded
+
+    return json_output