Update app.py

app.py

@@ -2,7 +2,6 @@ from transformers import GPT2LMHeadModel, GPT2Tokenizer
 from langchain import PromptTemplate
 from langchain.embeddings import HuggingFaceEmbeddings
 from langchain_community.vectorstores import FAISS
-from langchain_community.llms import CTransformers  # You might need to change this if GPT-2 isn't directly supported
 from langchain.chains import RetrievalQA
 import gradio as gr
 from huggingface_hub import hf_hub_download
@@ -40,19 +39,25 @@ Helpful answer:
     prompt = PromptTemplate(template=custom_prompt_template, input_variables=['context', 'question'])
     return prompt
 
-def retrieval_QA_chain(llm, tokenizer, prompt, db):
+def generate_answer(prompt_text, model, tokenizer):
     """
-    Create a RetrievalQA chain with the specified LLM, prompt, and vector store.
+    Generate an answer using the GPT-2 model and tokenizer.
     """
-    def generate_answer(query):
-        # Tokenize the input query
-        inputs = tokenizer.encode(query, return_tensors='pt')
-        # Generate response
-        outputs = llm.generate(inputs, max_length=512, temperature=0.5)
-        return tokenizer.decode(outputs[0], skip_special_tokens=True)
+    inputs = tokenizer.encode(prompt_text, return_tensors='pt')
+    outputs = model.generate(inputs, max_length=512, temperature=0.5)
+    return tokenizer.decode(outputs[0], skip_special_tokens=True)
 
+def retrieval_QA_chain(model, tokenizer, prompt, db):
+    """
+    Create a RetrievalQA chain with the specified LLM, prompt, and vector store.
+    """
+    def generate_answer_fn(query):
+        # Format the query with the prompt
+        formatted_prompt = prompt.format(context="Some context here", question=query)
+        return generate_answer(formatted_prompt, model, tokenizer)
+    
     qachain = RetrievalQA.from_chain_type(
-        llm=generate_answer,
+        llm=generate_answer_fn,
         chain_type="stuff",
         retriever=db.as_retriever(search_kwargs={'k': 2}),
         return_source_documents=True,
@@ -113,3 +118,4 @@ demo = gr.Interface(
 
 if __name__ == "__main__":
     demo.launch()
+