Update app.py

app.py

@@ -2,35 +2,52 @@ import gradio as gr
 from transformers import ViltProcessor, ViltForQuestionAnswering
 import torch
 
-# Load example images 
-torch.hub.download_url_to_file('http://images.cocodataset.org/val2017/000000039769.jpg', 'cats.jpg')
-
 # Load Vilt 
 vilt_processor = ViltProcessor.from_pretrained("dandelin/vilt-b32-finetuned-vqa")
 vilt_model = ViltForQuestionAnswering.from_pretrained("dandelin/vilt-b32-finetuned-vqa")
+
 def vilt_vqa(image, question):
-  # prepare inputs
   inputs = vilt_processor(image, question, return_tensors="pt")
-
-  # forward pass
   with torch.no_grad():
       outputs = vilt_model(**inputs)
-
   logits = outputs.logits
   idx = logits.argmax(-1).item()
   answer = vilt_model.config.id2label[idx]
   return answer
 
+# Load FLAN-T5
+t5_tokenizer = T5Tokenizer.from_pretrained("google/flan-t5-large")
+t5_model = T5ForConditionalGeneration.from_pretrained("google/flan-t5-large", device_map="auto")
+
+def flan_t5_complete_sentence(question, answer):
+  input_text = f"A question: {question} An incomplete answer: {answer}. Based on these, answer the question with a complete sentence without extra information."
+  print(input_text)
+  inputs = t5_tokenizer(input_text, return_tensors="pt")
+  outputs = t5_model.generate(**inputs, max_length=50)
+  result_sentence = t5_tokenizer.batch_decode(outputs, skip_special_tokens=True)
+  return result_sentence
+
+# Main function
+def vqa_main(image, question):
+    incomplete_answer = vilt_vqa(image, question)
+    complete_answer = flan_t5_complete_sentence(question, answer)
+    return complete_answer
+    
+# Home page text
+title = "Interactive demo: Multilingual VQA"
+description = "Demo for Multilingual VQA. Upload an image, type a question, click 'submit', or click one of the examples to load them."
+article = "article goes here"
+
+# Load example images 
+torch.hub.download_url_to_file('http://images.cocodataset.org/val2017/000000039769.jpg', 'cats.jpg')
+
+# Define home page variables 
 image = gr.inputs.Image(type="pil")
 question = gr.inputs.Textbox(label="Question")
 answer = gr.outputs.Textbox(label="Predicted answer")
 examples = [["cats.jpg", "What are the animals here called?"]]
 
-title = "Interactive demo: Multilingual VQA"
-description = "Demo for Multilingual VQA. Upload an image, type a question, click 'submit', or click one of the examples to load them."
-article = "article"
-
-interface = gr.Interface(fn=vilt_vqa, 
+interface = gr.Interface(fn=vqa_main, 
                          inputs=[image, question], 
                          outputs=answer, 
                          examples=examples,