Update app.py

app.py

@@ -15,77 +15,19 @@ ensemble_retriever = EnsembleRetriever(
     retrievers=[bm25_retriever, retriever], weights=[0.5, 0.5]
 )
 
-tokenizer = AutoTokenizer.from_pretrained("ShynBui/vie_qa", token=os.environ.get("HF_TOKEN"))
-model = AutoModelForQuestionAnswering.from_pretrained("ShynBui/vie_qa", token=os.environ.get("HF_TOKEN"))
-
-headers = {
-    "Accept": "application/json",
-    "Authorization": "Bearer " + os.environ.get("HF_TOKEN"),
-    "Content-Type": "application/json"
-}
-
-
-def query(payload):
-    response = requests.post(os.environ.get("API_URL"), headers=headers, json=payload)
-    return response.json()
-
-
-def greet(quote):
-    sources = []
-    answers = []
-    scores = []
-    ids = []
-
-    docs = ensemble_retriever.get_relevant_documents(quote)
-
-    for i in docs:
-        context = ViTokenizer.tokenize(i.page_content)
-        question = ViTokenizer.tokenize(quote)
-        print("source:", i.metadata['source'])
-        sources.append(i.metadata['source'])
-        output = query({
-            "inputs": {
-                "question": question,
-                "context": context[:256]
-            },
-        })
-        while "error" in output:
-            # print('fail')
-            time.sleep(1)
-            output = query({
-                "inputs": {
-                    "question": question,
-                    "context": context[:256]
-                },
-            })
-
-        answers.append(output['answer'])
-    return answers
 
 
 def greet2(quote):
-    answers = []
-    docs = ensemble_retriever.get_relevant_documents(quote)
-
-    return docs
-    
-    for i in docs:
-        context = ViTokenizer.tokenize(i.page_content)
-        question = ViTokenizer.tokenize(quote)
-
-        inputs = tokenizer(question, context, return_tensors="pt")
 
-        outputs = model(**inputs)
+    qa_chain = get_qachain(retriever=ensemble_retriever)
 
-        start_index = torch.argmax(outputs.start_logits)
-        end_index = torch.argmax(outputs.end_logits) + 1
+    prompt = os.environ['PROMPT']
 
-        answer = tokenizer.convert_tokens_to_string(
-            tokenizer.convert_ids_to_tokens(inputs["input_ids"][0][start_index:end_index]))
+    qa_chain.combine_documents_chain.llm_chain.prompt.messages[0].prompt.template = prompt
 
-        answers.append(answer)
+    llm_response = qa_chain(quote)
 
-    return answers
+    return llm_response['result']
 
 
 if __name__ == "__main__":