Update app.py

app.py

@@ -3,22 +3,30 @@ import gradio as gr
 from transformers import AutoModelForSeq2SeqLM, AutoTokenizer
 import torch
 import os
+import re
 
-# Load the model (flan-t5-base) and tokenizer
+# Load the model and tokenizer
 model_name = "google/flan-t5-base"
-hf_token = os.environ.get("HF_TOKEN")  # Ensure your token is securely set as a secret
+hf_token = os.environ.get("HF_TOKEN")  # Set as a secret in Hugging Face Space settings
 
 tokenizer = AutoTokenizer.from_pretrained(model_name, use_auth_token=hf_token)
 model = AutoModelForSeq2SeqLM.from_pretrained(model_name, use_auth_token=hf_token)
 
 # Move the model to CPU (or GPU if available)
-model.to("cpu")
+device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+model.to(device)
 
-# Function to generate the prompt for MT QE
+# Function to generate a clean prompt
 def generate_prompt(original, translation):
-    return f"### Task: Machine Translation Quality Estimation\n\nSource: {original}\nTranslation: {translation}\n\nScore (0-1):"
+    return (
+        f"Rate the quality of this translation from 0 (poor) to 1 (excellent). "
+        f"Only respond with a number.\n\n"
+        f"Source: {original}\n"
+        f"Translation: {translation}\n"
+        f"Score:"
+    )
 
-# Function to predict quality scores from the file
+# Main prediction function
 def predict_scores(file):
     df = pd.read_csv(file.name, sep="\t")
     scores = []
@@ -26,32 +34,35 @@ def predict_scores(file):
     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)
+        # Tokenize and send to model
+        inputs = tokenizer(prompt, return_tensors="pt").to(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 (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 in case of error
+        # Debug print (optional)
+        print("Response:", response)
+
+        # Extract numeric score using regex
+        match = re.search(r"\b([01](?:\.\d+)?)\b", response)
+        if match:
+            score_val = float(match.group(1))
+            score_val = max(0, min(score_val, 1))  # Clamp between 0 and 1
+        else:
+            score_val = -1  # fallback if model output is invalid
 
         scores.append(score_val)
 
     df["predicted_score"] = scores
     return df
 
-# Set up the Gradio interface
+# Gradio UI
 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 Google FLAN-T5-Base",
+    title="MT QE with FLAN-T5-Base",
+    description="Upload a dev.tsv file with columns: 'original' and 'translation'."
 )
 
-# Launch the Gradio interface
+# Launch app
 iface.launch()