Update app.py

app.py

@@ -82,24 +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 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")
+# Load Falcon-7B-Instruct model to rewrite answers
+gpt_tokenizer = AutoTokenizer.from_pretrained("tiiuae/falcon-7b-instruct")
+gpt_model = AutoModelForCausalLM.from_pretrained("tiiuae/falcon-7b-instruct")
 
 def rewrite_answer(question, short_answer):
-    prompt = f"Question: {question}\nShort Answer: {short_answer}\nFull sentence:"  # Few-shot style prompt
+    prompt = f"Rewrite the short answer into a natural sentence.\nQuestion: {question}\nAnswer: {short_answer}\nFull Sentence:"
     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.95,
+            top_p=0.9,
             temperature=0.8,
             pad_token_id=gpt_tokenizer.eos_token_id
         )
     rewritten = gpt_tokenizer.decode(outputs[0], skip_special_tokens=True)
-    return rewritten.split("Full sentence:")[-1].strip()
+    return rewritten.split("Full Sentence:")[-1].strip()
 
 def answer_question_from_image(image, question):
     if image is None or not question.strip():