from sentence_transformers import SentenceTransformer
import numpy as np
import logging, pickle,os
import faiss
import gradio as gr
INDEX_PATH = r"data/faiss_index.idx"
DOCUMENT_PATH = r"data/faiss_documents.pk1"
MODEL_NAME = 'all-MiniLM-L6-v2'
TOP_K_DEFAULT = 1
model = SentenceTransformer(MODEL_NAME)
def query_engine(query, top_k = TOP_K_DEFAULT):
    logging.info("Loading faiss index and documents")
    loaded_index = faiss.read_index(INDEX_PATH)
    with open(DOCUMENT_PATH,'rb') as f:
        docs = pickle.load(f)
    q_emb = model.encode([query],convert_to_numpy= True)
    q_emb = np.asarray(q_emb,dtype = np.float32)
    q_emb = np.ascontiguousarray(q_emb)
    faiss.normalize_L2(q_emb)
    D,I = loaded_index.search(q_emb, top_k)
    # print("FAISS score (inner product, larger is better)",D)
    # print("I is ",I)
    results = []
    for idx in I[0]:
        if idx < 0:
            continue
        results.append({
            'idx':int(idx),
            'instruction': docs[idx]['instruction'],
            'output': docs[idx]['output'],
            'text_for_embedding': docs[idx]['text_for_embedding']
        })
    return results[0]['output']

demo = gr.Interface(
    fn = query_engine,
    inputs = [gr.Textbox(label = "Question", lines = 4, max_lines = 8)],
    outputs = [gr.Textbox(label = "Answer", lines = 4, max_lines = 8)],
    title = "DSCM KNOWLEDGE BOT"
)

if __name__ == "__main__":
    demo.launch()