Update app.py

app.py

@@ -70,7 +70,7 @@ from fastapi.responses import RedirectResponse, FileResponse, JSONResponse
 import os
 import shutil
 from PIL import Image
-from transformers import ViltProcessor, ViltForQuestionAnswering, AutoTokenizer, AutoModelForCausalLM
+from transformers import ViltProcessor, ViltForQuestionAnswering, AutoTokenizer, AutoModelForSeq2SeqLM
 from gtts import gTTS
 import torch
 import tempfile
@@ -82,39 +82,26 @@ app = FastAPI()
 vqa_processor = ViltProcessor.from_pretrained("dandelin/vilt-b32-finetuned-vqa")
 vqa_model = ViltForQuestionAnswering.from_pretrained("dandelin/vilt-b32-finetuned-vqa")
 
-# Load GPT model to rewrite answers (Phi-1.5)
-gpt_tokenizer = AutoTokenizer.from_pretrained("microsoft/phi-1_5")
-gpt_model = AutoModelForCausalLM.from_pretrained("microsoft/phi-1_5")
+# Load FLAN-T5 model to rewrite answers (better for CPU)
+gpt_tokenizer = AutoTokenizer.from_pretrained("google/flan-t5-small")
+gpt_model = AutoModelForSeq2SeqLM.from_pretrained("google/flan-t5-small")
 
 def rewrite_answer(question, short_answer):
-    prompt = f"Write a full sentence to answer this:\nQ: {question}\nA: {short_answer}\nFull sentence:"
+    prompt = f"Write a full sentence that answers the question '{question}' using this answer: {short_answer}."
     inputs = gpt_tokenizer(prompt, return_tensors="pt")
     with torch.no_grad():
         outputs = gpt_model.generate(
             **inputs,
             max_new_tokens=50,
             do_sample=True,
-            top_k=40,
             top_p=0.9,
-            temperature=0.6,
-            pad_token_id=gpt_tokenizer.eos_token_id
+            temperature=0.7,
+            pad_token_id=gpt_tokenizer.pad_token_id
         )
 
-    generated = gpt_tokenizer.decode(outputs[0], skip_special_tokens=True)
-
-    # Try to isolate the answer portion
-    if "Full sentence:" in generated:
-        rewritten = generated.split("Full sentence:")[-1].strip()
-    else:
-        rewritten = generated.strip()
-
-    # Fallback to basic templating if model fails
-    if not rewritten or len(rewritten.split()) < 3:
-        rewritten = f"The answer to the question '{question}' is: {short_answer}."
-
+    rewritten = gpt_tokenizer.decode(outputs[0], skip_special_tokens=True)
     return rewritten
 
-
 def answer_question_from_image(image, question):
     if image is None or not question.strip():
         return "Please upload an image and ask a question.", None
@@ -125,7 +112,7 @@ def answer_question_from_image(image, question):
     predicted_id = outputs.logits.argmax(-1).item()
     short_answer = vqa_model.config.id2label[predicted_id]
 
-    # Rewrite short answer to full sentence with Phi-1.5
+    # Rewrite short answer to full sentence
     full_answer = rewrite_answer(question, short_answer)
 
     try: