import gradio as gr
import os 
from huggingface_hub import login

from torch import float16

login(os.getenv('Token'))
 

from transformers import AutoModelForCausalLM
from transformers import AutoTokenizer 

# Load the model and tokenizer
model_name ="meta-llama/Llama-2-7b-chat-hf" #"meta-llama/Llama-3.1-8B"  # Replace with your desired Hugging Face model
model = AutoModelForCausalLM.from_pretrained(model_name, torch_dtype=float16)
tokenizer = AutoTokenizer.from_pretrained(model_name)


def respond(
    message,
    history: list[tuple[str, str]],
    system_message,
    max_tokens,
    temperature,
    top_p,
):
    prompt = ""
    if system_message:
        prompt += system_message + "\n"
    for user_message, model_response in history:
        prompt += f"User: {user_message}\nAssistant: {model_response}\n"
    prompt += f"User: {message}\nAssistant: "

    print("Prompt", prompt)

    # Tokenize the prompt
    inputs = tokenizer(prompt, return_tensors="pt")
    
    print("Input", inputs)

    # Generate text
    outputs = model.generate(
        **inputs, 
        max_length=max_tokens or 512, #+ len(inputs["input_ids"][0]), 
        do_sample=True, 
        temperature=temperature or 1.0, 
        top_p=top_p or 0.9
    )

    print("Output", outputs)

    # Decode the generated text
    generated_text = tokenizer.decode(outputs[0], skip_special_tokens=True) 

    print('the response',generated_text)
    # Extract the assistant's response from the generated text
    response = generated_text.split("Assistant: ")[-1]

    print('the response 2',response)

    return response


"""
For information on how to customize the ChatInterface, peruse the gradio docs: https://www.gradio.app/docs/chatinterface
"""
demo = gr.ChatInterface(
    respond,
    type="messages",
    additional_inputs=[
        gr.Textbox(value="You are a friendly Chatbot.", label="System message"),
        gr.Slider(minimum=1, maximum=2048, value=512, step=1, label="Max new tokens"),
        gr.Slider(minimum=0.1, maximum=4.0, value=0.7, step=0.1, label="Temperature"),
        gr.Slider(
            minimum=0.1,
            maximum=1.0,
            value=0.95,
            step=0.05,
            label="Top-p (nucleus sampling)",
        ),
    ]
)


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