from sentence_transformers import SentenceTransformer

import torch
import gradio as gr
from huggingface_hub import InferenceClient

with open("hindu_yuva_knowledge_base.txt", "r", encoding="utf-8") as file:
    yuva_knowledge = file.read()

def preprocess_text(text):
    cleaned_text = text.strip()
    chunks = cleaned_text.split("\n")
    cleaned_chunks = []
    for chunk in chunks:
        chunk = chunk.strip()
        if chunk != "":
            cleaned_chunks.append(chunk)
    return cleaned_chunks

cleaned_chunks = preprocess_text(yuva_knowledge)


model = SentenceTransformer('all-MiniLM-L6-v2')

def create_embeddings(text_chunks):
    chunk_embeddings = model.encode(text_chunks, convert_to_tensor=True)
    return chunk_embeddings

chunk_embeddings = create_embeddings(cleaned_chunks)

def get_top_chunks(query, chunk_embeddings, text_chunks):
    query_embedding = model.encode(query, convert_to_tensor=True)
    query_embedding_normalized = query_embedding / query_embedding.norm()
    chunk_embeddings_normalized = chunk_embeddings / chunk_embeddings.norm(dim=1, keepdim=True)
    similarities = torch.matmul(chunk_embeddings_normalized, query_embedding_normalized)
    top_indices = torch.topk(similarities, k=3).indices
    top_chunks = []
    for index in top_indices:
        top_chunks.append(text_chunks[index])
    return top_chunks

client = InferenceClient("Qwen/Qwen2.5-7B-Instruct")

def respond(message, history):
    top_chunks = get_top_chunks(message, chunk_embeddings, cleaned_chunks)
    context = "\n\n".join(top_chunks)
    messages = [{"role": "system", "content": f"Always greet the user first by asking 'Hello, how can I help you?'. You are an assistant answering users' questions about Hindu YUVA. You just need to pull information from the website to answer basic questions. \n{context}"}]
    for turn in history:
        if turn["role"] == "user":
            messages.append({"role": "user", "content": turn["content"]})
        elif turn["role"] == "assistant":
            messages.append({"role": "assistant", "content": turn["content"]})
    messages.append({"role": "user", "content": message})

    response = ""
    for msg in client.chat_completion(messages, stream=True):
        token = msg.choices[0].delta.content
        if token is not None:
            response += token
            yield response

chatbot = gr.ChatInterface(respond)
chatbot.launch(debug=True)