Update app.py

app.py

@@ -10,10 +10,12 @@ embed_model = SentenceTransformer(
     "sentence-transformers/all-MiniLM-L6-v2"
 )
 
-generator = pipeline(
-    "text2text-generation",
-    model="google/flan-t5-base"
-)
+from transformers import AutoTokenizer, AutoModelForSeq2SeqLM
+
+gen_model_name="google/flan-t5-base"
+
+gen_tokenizer=AutoTokenizer.from_pretrained(gen_model_name)
+gen_model=AutoModelForSeq2SeqLM.from_pretrained(gen_model_name)
 
 
 chunks_store=[]
@@ -62,7 +64,7 @@ def ask_question(question):
     )
 
     prompt=f"""
-Use context to answer.
+Answer the question only using the context.
 
 Context:
 {retrieved}
@@ -73,13 +75,25 @@ Question:
 Answer:
 """
 
-    answer=generator(
+    inputs=gen_tokenizer(
         prompt,
-        max_length=150
-    )[0]["generated_text"]
+        return_tensors="pt",
+        truncation=True,
+        max_length=512
+    )
 
-    return answer,retrieved
+    outputs=gen_model.generate(
+        **inputs,
+        max_new_tokens=120,
+        num_beams=4
+    )
+
+    answer=gen_tokenizer.decode(
+        outputs[0],
+        skip_special_tokens=True
+    )
 
+    return answer,retrieved
 
 
 with gr.Blocks() as demo: