Update app.py

app.py

@@ -19,9 +19,9 @@ client = chromadb.Client()
 collection = client.get_or_create_collection("patents")
 collection.add(documents=docs,ids=ids)
 
-#def text_embedding(input):
-    #model = SentenceTransformer('sentence-transformers/all-MiniLM-L6-v2')
-    #return model.encode(input)
+def text_embedding(input):
+    model = SentenceTransformer('sentence-transformers/all-MiniLM-L6-v2')
+    return model.encode(input)
 
 def gen_context(query):
     vector = text_embedding(query).tolist()    
@@ -29,13 +29,20 @@ def gen_context(query):
     res = "\n".join(str(item) for item in results['documents'][0])
     return res
 
-def chat_completion(user_prompt):
-    length = 1000
-    system_prompt = """\You are a helpful AI assistant that can answer questions on the patents dataset. Answer based on the context provided.If you cannot find the correct answer, say I don't know. Be concise and just include the response"""
-    final_prompt = f"""<s>[INST]<<SYS>>{system_prompt}<</SYS>>{user_prompt}[/INST]"""
+
+def chat_completion(query):
+    
+    length = 1000    
+    context = gen_context(query)
+    
+    user_prompt = f"""Based on the context:{context}Answer the below query:{query}"""
+    system_prompt = """You are a helpful AI assistant that can answer questions on the patents dataset. Answer based on the context provided.If you cannot find the correct answer, say I don't know. Be concise and just include the response"""
+    final_prompt = f"""<s>[INST]<<SYS>>{system_prompt}<</SYS>>{user_prompt}[/INST]""" 
+    
     
     return client.text_generation(prompt=final_prompt,max_new_tokens = length).strip()
 
+
 client = InferenceClient(model = "mistralai/Mixtral-8x7B-Instruct-v0.1")
 
 demo = gr.Interface(fn=chat_completion,