Update app.py

app.py

@@ -65,15 +65,13 @@ app = gr.mount_gradio_app(app, gui, path="/")
 @app.get("/")
 def home():
     return RedirectResponse(url="/") """
-from fastapi import FastAPI, UploadFile, Form
-from fastapi.responses import RedirectResponse, FileResponse, JSONResponse
-import os
-import shutil
+from fastapi import FastAPI
+from fastapi.responses import RedirectResponse
+import tempfile
 from PIL import Image
+import torch
 from transformers import ViltProcessor, ViltForQuestionAnswering, AutoTokenizer, AutoModelForCausalLM
 from gtts import gTTS
-import torch
-import tempfile
 import gradio as gr
 
 app = FastAPI()
@@ -86,19 +84,30 @@ vqa_model = ViltForQuestionAnswering.from_pretrained("dandelin/vilt-b32-finetune
 gpt_tokenizer = AutoTokenizer.from_pretrained("distilgpt2")
 gpt_model = AutoModelForCausalLM.from_pretrained("distilgpt2")
 
-
-def rewrite_answer(question):
-    prompt = f"{question}\nAnswer with a full sentence:"
+def rewrite_answer(question, short_answer):
+    prompt = (
+        f"Question: {question}\n"
+        f"Short Answer: {short_answer}\n"
+        f"Now write a full sentence answering the question:"
+    )
     inputs = gpt_tokenizer(prompt, return_tensors="pt")
     with torch.no_grad():
         outputs = gpt_model.generate(
             **inputs,
-            max_new_tokens=40,
-            do_sample=False,
+            max_new_tokens=50,
+            do_sample=True,
+            top_p=0.9,
+            temperature=0.7,
             pad_token_id=gpt_tokenizer.eos_token_id
         )
+
     generated = gpt_tokenizer.decode(outputs[0], skip_special_tokens=True)
-    rewritten = generated.split(":")[-1].strip()
+
+    if "Now write a full sentence answering the question:" in generated:
+        rewritten = generated.split("Now write a full sentence answering the question:")[-1].strip()
+    else:
+        rewritten = generated.strip()
+
     return rewritten
 
 def answer_question_from_image(image, question):
@@ -111,16 +120,17 @@ 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 GPT-Neo
-    full_answer = rewrite_answer(f"Question: {question}\nAnswer: {short_answer}")
+    # Rewrite to human-like sentence
+    full_answer = rewrite_answer(question, short_answer)
 
+    # Convert to speech
     try:
         tts = gTTS(text=full_answer)
         with tempfile.NamedTemporaryFile(delete=False, suffix=".mp3") as tmp:
             tts.save(tmp.name)
             audio_path = tmp.name
     except Exception as e:
-        return f"Answer: {full_answer}\n\n⚠️ Audio generation error: {e}", None
+        return f"{full_answer}\n\n⚠️ Audio generation error: {e}", None
 
     return full_answer, audio_path
 
@@ -128,6 +138,7 @@ def process_image_question(image: Image.Image, question: str):
     answer, audio_path = answer_question_from_image(image, question)
     return answer, audio_path
 
+# Gradio UI
 gui = gr.Interface(
     fn=process_image_question,
     inputs=[
@@ -139,11 +150,12 @@ gui = gr.Interface(
         gr.Audio(label="Answer (Audio)", type="filepath")
     ],
     title="🧠 Image QA with Voice",
-    description="Upload an image and ask a question. You'll get a full-sentence spoken answer."
+    description="Upload an image and ask a question. You'll get a human-like spoken answer."
 )
 
+# Mount on FastAPI
 app = gr.mount_gradio_app(app, gui, path="/")
 
 @app.get("/")
 def home():
-    return RedirectResponse(url="/")
+    return RedirectResponse(url="/")