Update utils.py

utils.py

@@ -6,6 +6,7 @@ from langchain.vectorstores import Chroma
 from transformers import AutoTokenizer, AutoModelForCausalLM, pipeline
 from langchain_community.llms import HuggingFacePipeline
 from langchain.chains.question_answering import load_qa_chain
+from sentence_transformers import SentenceTransformer
 
 # Load and process documents
 dir = "data"
@@ -31,7 +32,7 @@ vectordb = Chroma.from_documents(docs, embeddings, persist_directory=persist_dir
 vectordb.persist()
 new_db = Chroma(persist_directory=persist_directory, embedding_function=embeddings)
 
-def get_similar_docs(query, k=1, score=False):
+def get_similar_docs(query, k=2, score=False):
     if score:
         similar_docs = new_db.similarity_search_with_score(query, k=k)
     else:
@@ -39,20 +40,25 @@ def get_similar_docs(query, k=1, score=False):
     return similar_docs
 
 # Load LLM model from Hugging Face
-model_name = "HuggingFaceH4/zephyr-7b-beta"
 
-llm = HuggingFacePipeline.from_model_id(
-    model_id=model_name,
+model_name = "HuggingFaceH4/zephyr-7b-beta"
+model = AutoModelForCausalLM.from_pretrained(model_name)
+tokenizer = AutoTokenizer.from_pretrained(model_name)
+    
+text_generation_pipeline = pipeline(
+    model=model,
+    tokenizer=tokenizer,
     task="text-generation",
-    pipeline_kwargs={
-        "max_new_tokens": 400,
-        "top_k": 50,
-        "temperature": 0.2,
-        "do_sample": True,
-        "repetition_penalty": 1.1,
-    }
+    temperature=0.2,
+    do_sample=True,
+    repetition_penalty=1.1,
+    return_full_text=True,
+    max_new_tokens=400,
+    inference= True,
 )
 
+llm = HuggingFacePipeline(pipeline=text_generation_pipeline)
+
 chain = load_qa_chain(llm, chain_type="stuff")
 
 def get_helpful_answer(text):