Vector_db / app.py
ShynBui's picture
Update app.py
038c95d verified
raw
history blame
No virus
904 Bytes
import gradio as gr
import os
from langchain.retrievers import EnsembleRetriever
from utils import *
import requests
from pyvi import ViTokenizer, ViPosTagger
import time
from transformers import AutoTokenizer, AutoModelForQuestionAnswering
import torch
retriever = load_the_embedding_retrieve(is_ready=False, k=5)
bm25_retriever = load_the_bm25_retrieve(k=5)
ensemble_retriever = EnsembleRetriever(
retrievers=[bm25_retriever, retriever], weights=[0.5, 0.5]
)
def greet2(quote):
qa_chain = get_qachain(retriever=ensemble_retriever)
prompt = os.environ['PROMPT']
qa_chain.combine_documents_chain.llm_chain.prompt.messages[0].prompt.template = prompt
llm_response = qa_chain(quote)
return llm_response['result']
if __name__ == "__main__":
quote = "Địa chỉ nhà trường?"
iface = gr.Interface(fn=greet2, inputs="text", outputs="text")
iface.launch()