Update app.py

app.py

@@ -1,51 +1,57 @@
 import pandas as pd
 import gradio as gr
-from transformers import AutoModelForCausalLM, AutoTokenizer, BitsAndBytesConfig
+from transformers import AutoModelForSeq2SeqLM, AutoTokenizer
 import torch
 import os
 
-# Load model
-bnb_config = BitsAndBytesConfig(load_in_4bit=True)
-model_name = "meta-llama/Meta-Llama-3-8B-Instruct"
-hf_token = os.environ.get("HF_TOKEN")
+# Load the model (flan-t5-base) and tokenizer
+model_name = "google/flan-t5-base"
+hf_token = os.environ.get("HF_TOKEN")  # Ensure your token is securely set as a secret
 
 tokenizer = AutoTokenizer.from_pretrained(model_name, use_auth_token=hf_token)
-model = AutoModelForCausalLM.from_pretrained(model_name, use_auth_token=hf_token)
-model.to("cpu")  # move model to CPU (if not using GPU Space)
+model = AutoModelForSeq2SeqLM.from_pretrained(model_name, use_auth_token=hf_token)
 
-# Prompt generator
+# Move the model to CPU (or GPU if available)
+model.to("cpu")
+
+# Function to generate the prompt for MT QE
 def generate_prompt(original, translation):
     return f"### Task: Machine Translation Quality Estimation\n\nSource: {original}\nTranslation: {translation}\n\nScore (0-1):"
 
-# Prediction function
+# Function to predict quality scores from the file
 def predict_scores(file):
     df = pd.read_csv(file.name, sep="\t")
     scores = []
 
     for _, row in df.iterrows():
         prompt = generate_prompt(row["original"], row["translation"])
+
+        # Tokenize and generate outputs
         inputs = tokenizer(prompt, return_tensors="pt").to(model.device)
         outputs = model.generate(**inputs, max_new_tokens=10)
+
+        # Decode and extract the score from the response
         response = tokenizer.decode(outputs[0], skip_special_tokens=True)
 
-        # Extract float value (naive way)
+        # Extract float value (simple way to extract score from response)
         score = response.split("Score")[-1].strip()
         try:
             score_val = float(score.split()[0])
         except:
-            score_val = -1  # fallback
+            score_val = -1  # Fallback in case of error
 
         scores.append(score_val)
 
     df["predicted_score"] = scores
     return df
 
-# Gradio app
+# Set up the Gradio interface
 iface = gr.Interface(
     fn=predict_scores,
     inputs=gr.File(label="Upload dev.tsv"),
     outputs=gr.Dataframe(label="QE Output with Predicted Score"),
-    title="MT QE with LLaMA-3.2-3B-Instruct"
+    title="MT QE with Google FLAN-T5-Base",
 )
 
+# Launch the Gradio interface
 iface.launch()