from langchain.vectorstores import FAISS
from langchain.embeddings import SentenceTransformerEmbeddings
import gradio as gr
import reranking
from extract_keywords import init_keyword_extractor, extract_keywords

embeddings = SentenceTransformerEmbeddings(model_name="multi-qa-MiniLM-L6-cos-v1")
db = FAISS.load_local('faiss_qa', embeddings)
init_keyword_extractor()

def main(query):
    query = query.lower()
    query_keywords = set(extract_keywords(query))
    result_docs = db.similarity_search_with_score(query, k=20)

    if len(query_keywords) > 0:
        result_docs = filter(lambda doc: len(set(extract_keywords(doc[0].page_content)).intersection(query_keywords)) > 0, result_docs)

    if len(result_docs) == 0:
        return 'Ответ не найден', 0, ''
      
    if len(result_docs) == 1:
        score, index = 0, 0
    else:
        sentences = [doc[0].page_content for doc in result_docs]
        #print('----------------------------------------------------------------')
        #for doc in result_docs:
        #    print(doc[0].metadata['articleId'], ' | ', doc[0].page_content, ' | ', doc[0].metadata['answer'])
        score, index = reranking.search(query, sentences)

    return result_docs[index][0].metadata['answer'], score, result_docs[index][0].page_content

demo = gr.Interface(fn=main, inputs="text", outputs=[
  gr.Textbox(label="Ответ, который будет показан клиенту"),
  gr.Textbox(label="Score"),
  gr.Textbox(label="Вопрос, по которому был найден ответ"),
])

demo.launch()