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, AutoModelForSeq2SeqLM
+from transformers import ViltProcessor, ViltForQuestionAnswering, AutoTokenizer, AutoModelForCausalLM
 from gtts import gTTS
 import torch
 import tempfile
@@ -82,25 +82,24 @@ app = FastAPI()
 vqa_processor = ViltProcessor.from_pretrained("dandelin/vilt-b32-finetuned-vqa")
 vqa_model = ViltForQuestionAnswering.from_pretrained("dandelin/vilt-b32-finetuned-vqa")
 
-# 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")
+# Load Falcon-RW-1B model to rewrite answers
+gpt_tokenizer = AutoTokenizer.from_pretrained("tiiuae/falcon-rw-1b")
+gpt_model = AutoModelForCausalLM.from_pretrained("tiiuae/falcon-rw-1b")
 
 def rewrite_answer(question, short_answer):
-    prompt = f"Write a full sentence that answers the question '{question}' using this answer: {short_answer}."
+    prompt = f"Question: {question}\nShort Answer: {short_answer}\nFull sentence:"  # Few-shot style prompt
     inputs = gpt_tokenizer(prompt, return_tensors="pt")
     with torch.no_grad():
         outputs = gpt_model.generate(
             **inputs,
             max_new_tokens=50,
             do_sample=True,
-            top_p=0.9,
-            temperature=0.7,
-            pad_token_id=gpt_tokenizer.pad_token_id
+            top_p=0.95,
+            temperature=0.8,
+            pad_token_id=gpt_tokenizer.eos_token_id
         )
-
     rewritten = gpt_tokenizer.decode(outputs[0], skip_special_tokens=True)
-    return rewritten
+    return rewritten.split("Full sentence:")[-1].strip()
 
 def answer_question_from_image(image, question):
     if image is None or not question.strip():
@@ -147,4 +146,4 @@ app = gr.mount_gradio_app(app, gui, path="/")
 
 @app.get("/")
 def home():
-    return RedirectResponse(url="/")
+    return RedirectResponse(url="/")