Update app.py

app.py

@@ -3,11 +3,25 @@ from langchain import PromptTemplate
 from langchain.embeddings import HuggingFaceEmbeddings
 from langchain_community.vectorstores import FAISS
 from langchain.chains import RetrievalQA
+from langchain.chains.llm import LLMChain
+from langchain.chains.question_answering import load_qa_chain
 import gradio as gr
-from huggingface_hub import hf_hub_download
 
 DB_FAISS_PATH = "vectorstores/db_faiss"
 
+class GPT2LLM:
+    """
+    A custom class to wrap the GPT-2 model and tokenizer to be used with LangChain.
+    """
+    def __init__(self, model, tokenizer):
+        self.model = model
+        self.tokenizer = tokenizer
+
+    def __call__(self, prompt_text, max_length=512):
+        inputs = self.tokenizer.encode(prompt_text, return_tensors='pt')
+        outputs = self.model.generate(inputs, max_length=max_length, temperature=0.5)
+        return self.tokenizer.decode(outputs[0], skip_special_tokens=True)
+
 def load_llm():
     """
     Load the GPT-2 model for the language model.
@@ -18,10 +32,10 @@ def load_llm():
         model = GPT2LMHeadModel.from_pretrained(model_name)
         tokenizer = GPT2Tokenizer.from_pretrained(model_name)
         print("Model and tokenizer successfully loaded!")
-        return model, tokenizer
+        return GPT2LLM(model, tokenizer)
     except Exception as e:
         print(f"An error occurred while loading the model: {e}")
-        return None, None
+        return None
 
 def set_custom_prompt():
     """
@@ -39,29 +53,16 @@ Helpful answer:
     prompt = PromptTemplate(template=custom_prompt_template, input_variables=['context', 'question'])
     return prompt
 
-def generate_answer(prompt_text, model, tokenizer):
-    """
-    Generate an answer using the GPT-2 model and tokenizer.
-    """
-    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):
+def retrieval_QA_chain(llm, 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)
-    
+    llm_chain = LLMChain(llm=llm, prompt=prompt)
     qachain = RetrievalQA.from_chain_type(
-        llm=generate_answer_fn,
+        llm_chain=llm_chain,
         chain_type="stuff",
         retriever=db.as_retriever(search_kwargs={'k': 2}),
-        return_source_documents=True,
-        chain_type_kwargs={'prompt': prompt}
+        return_source_documents=True
     )
     return qachain
 
@@ -71,10 +72,10 @@ def qa_bot():
     """
     embeddings = HuggingFaceEmbeddings(model_name='sentence-transformers/all-miniLM-L6-V2', model_kwargs={'device': 'cpu'})
     db = FAISS.load_local(DB_FAISS_PATH, embeddings, allow_dangerous_deserialization=True)
-    model, tokenizer = load_llm()
+    llm = load_llm()
     qa_prompt = set_custom_prompt()
-    if model and tokenizer:
-        qa = retrieval_QA_chain(model, tokenizer, qa_prompt, db)
+    if llm:
+        qa = retrieval_QA_chain(llm, qa_prompt, db)
     else:
         qa = None
     return qa