import torch
from transformers import AutoModelForCausalLM, AutoTokenizer, BitsAndBytesConfig
from sentence_transformers import SentenceTransformer
from llama_index.core import Settings, VectorStoreIndex, SimpleDirectoryReader, PromptTemplate
from llama_index.llms.huggingface import HuggingFaceLLM
from llama_index.embeddings.huggingface import HuggingFaceEmbedding

documents = SimpleDirectoryReader("../Datasets/Merged").load_data()

def completion_to_prompt(completion):
    return f"<|system|>\n</s>\n<|user|>\n{completion}</s>\n<|assistant|>\n"

def messages_to_prompt(messages):
    prompt = """You are an expert in blockchain technologies and blockchain patterns.\n
    Your answer should be returned in the following format/ \n
    Summary: \n
    Context: \n
    Problem: \n
    Forces: \n
    Solution: \n
    Benefits: \n
    Drawbacks: \n
    Related patterns: \n
    Known uses: \n

    The context will be provided and you should base your answer on it.
    """
    for message in messages:
        if message.role == "system":
            prompt += f"<|system|>\n{message.content}</s>\n"
        elif message.role == "user":
            prompt += f"<|user|>\n{message.content}</s>\n"
        elif message.role == "assistant":
            prompt += f"<|assistant|>\n{message.content}</s>\n"
    if not prompt.startswith("<|system|>\n"):
        prompt = "<|system|>\n</s>\n" + prompt

    prompt = prompt + "<|assistant|>\n"

    return prompt
)

documents = SimpleDirectoryReader("../Datasets/Merged").load_data()

def completion_to_prompt(completion):
    return f"<|system|>\n</s>\n<|user|>\n{completion}</s>\n<|assistant|>\n"

# Transform a list of chat messages into zephyr-specific input
def messages_to_prompt(messages):
    prompt = ""
    for message in messages:
        if message.role == "system":
            prompt += f"<|system|>\n{message.content}</s>\n"
        elif message.role == "user":
            prompt += f"<|user|>\n{message.content}</s>\n"
        elif message.role == "assistant":
            prompt += f"<|assistant|>\n{message.content}</s>\n"

    # ensure we start with a system prompt, insert blank if needed
    if not prompt.startswith("<|system|>\n"):
        prompt = "<|system|>\n</s>\n" + prompt

    # add final assistant prompt
    prompt = prompt + "<|assistant|>\n"

    return prompt

Settings.llm = HuggingFaceLLM(
    model_name="./zephyr-7b-beta",
    tokenizer_name="./zephyr-7b-beta",
    context_window = 3900,
    max_new_tokens = 256,
    generate_kwargs={"temperature": 0, "top_k": 30, "top_p": 0.95, "do_sample": True},
    messages_to_prompt=messages_to_prompt,
    completion_to_prompt=completion_to_prompt,
    device_map="auto",
)

Settings.embed_model = HuggingFaceEmbedding(
    model_name="./BAAI/bge-small-en-v1.5"
)

index = VectorStoreIndex.from_documents(documents)

query_engine = index.as_query_engine()
response = query_engine.query("What are Zero-Knowledge Proofs ?")

print(response)