__import__('pysqlite3')
import sys
sys.modules['sqlite3'] = sys.modules.pop('pysqlite3')

import os
import gradio as gr
import chromadb
from sentence_transformers import SentenceTransformer
import pandas as pd
import numpy as np

from chromadb.utils import embedding_functions
from huggingface_hub import InferenceClient

dfs = pd.read_csv('Patents.csv')
ids= [str(x) for x in dfs.index.tolist()]
docs = dfs['text'].tolist()
client = chromadb.Client()
collection = client.get_or_create_collection("patents")
collection.add(documents=docs,ids=ids)

##

def text_embedding(input):
    model = SentenceTransformer('sentence-transformers/all-MiniLM-L6-v2')
    return model.encode(input)

def gen_context(query):
    vector = text_embedding(query).tolist()    
    results = collection.query(query_embeddings=vector,n_results=15,include=["documents"])    
    res = "\n".join(str(item) for item in results['documents'][0])
    return res


def chat_completion(query):
    
    length = 1000    
    context = gen_context(query)
    
    user_prompt = f"""Based on the context:{context}Answer the below query:{query}"""
    system_prompt = """You are a helpful AI assistant that can answer questions on the patents dataset. Answer based on the context provided.If you cannot find the correct answer, say I don't know. Be concise and just include the response"""
    final_prompt = f"""<s>[INST]<<SYS>>{system_prompt}<</SYS>>{user_prompt}[/INST]""" 
    
    
    return client.text_generation(prompt=final_prompt,max_new_tokens = length).strip()


client = InferenceClient(model = "mistralai/Mixtral-8x7B-Instruct-v0.1")

demo = gr.Interface(fn=chat_completion, 
                    inputs=[gr.Textbox(label="Query", lines=2)],
                    outputs=[gr.Textbox(label="Result", lines=16)],
                    title="Chat on Patents Data")

demo.queue().launch(share=True)