import transformers
import torch
from transformers import pipeline
import accelerate

from transformers import AutoTokenizer, AutoModelForCausalLM

model_id = "pankaj9075rawat/DevsDoCode_LLama-3-8b-Uncensored"

# Load the tokenizer and model
tokenizer = AutoTokenizer.from_pretrained(model_id)
model = AutoModelForCausalLM.from_pretrained(
    model_id,
    torch_dtype=torch.bfloat16,
    device_map="auto",
)

pipeline = pipeline(
    "text-generation",
    model=model,
    tokenizer=tokenizer,
    # model_kwargs={"torch_dtype": torch.bfloat16},
    # device="cuda",
    # device_map="auto",
    # token=access_token
)

terminators = [
    pipeline.tokenizer.eos_token_id,
    pipeline.tokenizer.convert_tokens_to_ids("<|eot_id|>")
]

def get_response(
          query, message_history=[], max_tokens=128, temperature=1.1, top_p=0.9
      ):
    user_prompt = message_history + [{"role": "user", "content": query}]
    prompt = pipeline.tokenizer.apply_chat_template(
        user_prompt, tokenize=False, add_generation_prompt=True
    )
    # print("prompt before coversion: ", user_prompt)
    # print("prompt after conversion: ", prompt)
    outputs = pipeline(
        prompt,
        max_new_tokens=max_tokens,
        eos_token_id=terminators,
        do_sample=True,
        temperature=temperature,
        top_p=top_p,
    )
    response = outputs[0]["generated_text"][len(prompt):]
    return response, user_prompt + [{"role": "assistant", "content": response}]

convers = [{"role": "system", "content": system_instruction}]


def chat():
    global convers 
    response, convers = get_init_AI_response(convers)
    print("response:", response)

    while True:
        user_input = input("enter chat")
        if user_input.lower() in ["exit", "quit"]:
            return {"response": "Exiting the chatbot. Goodbye!"}

        response, convers = get_response(user_input, convers)
        print("response:", response)

chat()