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=3)
bm25_retriever = load_the_bm25_retrieve(k=3)

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()