fixing llm response

src/streamlit_app.py

@@ -40,13 +40,14 @@ def load_retriever():
 # Load a lightweight model via HuggingFace pipeline
 @st.cache_resource
 def load_llm():
-    pipe = pipeline("text-generation", model="google/flan-t5-small", max_new_tokens=256)
+    # pipe = pipeline("text-generation", model="google/flan-t5-small", max_new_tokens=256)
     # load the tokenizer and model on cpu/gpu
 
     model_name = "meta-llama/Llama-2-7b-chat-hf"
     tokenizer = AutoTokenizer.from_pretrained(model_name)
     model = AutoModelForCausalLM.from_pretrained(model_name, torch_dtype="auto", device_map="auto")
-    # pipe = pipeline("text-generation", model=model, tokenizer=tokenizer, max_new_tokens=256)
+    pipe = pipeline("text-generation", model=model, tokenizer=tokenizer, max_new_tokens=256)
+    
     return HuggingFacePipeline(pipeline=pipe)
 
 # Setup RAG Chain
@@ -56,10 +57,10 @@ def setup_qa():
     retriever = load_retriever()
     llm = load_llm()
     question_answer_chain = create_stuff_documents_chain(llm,prompt)
-    chain = create_retrieval_chain(retriever, question_answer_chain)
+    # chain = create_retrieval_chain(retriever, question_answer_chain)
 
-    # qa_chain = RetrievalQA.from_chain_type(llm=llm, retriever=retriever)
-    return chain
+    qa_chain = RetrievalQA.from_chain_type(llm=llm, retriever=retriever, chain_type="stuff")
+    return qa_chain
 
 
 # Streamlit App UI