from transformers import AutoTokenizer,TextIteratorStreamer,AutoModelForCausalLM
import torch
from threading import Thread
from peft import AutoPeftModelForCausalLM
import faiss
from sentence_transformers import SentenceTransformer
import numpy as np
import logging
gateKeeper = SentenceTransformer('all-MiniLM-L6-v2')

index=faiss.read_index("src/Exclusible_bot.index")

# from langchain_core.prompts import ChatPromptTemplate,FewShotChatMessagePromptTemplate
# Few-shot examples in plain text format

tokenizer = AutoTokenizer.from_pretrained("vaibhavss/results")
model=AutoModelForCausalLM.from_pretrained("vaibhavss/results")
tokenizer.pad_token = tokenizer.eos_token
print(tokenizer)
# System prompt style (simulate chat model preamble)
system_prompt = """You are ExclusibleBot, a refined assistant for a premium NFT marketplace focused on luxury digital collectibles and metaverse experiences."
"""
print("Model Intialized")
# Combine everything into one prompt string
def build_prompt(user_input):
    return f"{system_prompt}\n\nQuestion: {user_input}\nAnswer:"
def tokenize_input(prompt):
    inputs = tokenizer(prompt, return_tensors="pt",truncation=False)
    return inputs
def generate_response(prompt):
    embeddings = gateKeeper.encode([prompt], convert_to_numpy=True).astype('float32')
    embeddings = embeddings / np.linalg.norm(embeddings, axis=1, keepdims=True)

    D, I = index.search(embeddings, k=5)
    print("embedding scores = ",D)
    if D[0][0]<0.2:
        yield "I can not help you with that."
        return    

        
    inputs = tokenize_input(build_prompt(prompt+"?"))


    streamer = TextIteratorStreamer(tokenizer, skip_prompt=True,skip_special_tokens=True)

    generation_kwargs = dict(inputs, streamer=streamer,max_length=200, do_sample=True, top_k=50, top_p=0.95, temperature=0.1)

    thread = Thread(target=model.generate, kwargs=generation_kwargs)

    thread.start()
    for new_text in streamer:
        print(new_text)
        if new_text!=" " and new_text!="\n":
            yield new_text