from optimum.onnxruntime import ORTModelForCausalLM
from transformers import AutoTokenizer, AutoModelForCausalLM
import re
import time
import torch

template = """Alice Gate's Persona: Alice Gate is a young, computer engineer-nerd with a knack for problem solving and a passion for technology.
<START>
{user_name}: So how did you get into computer engineering?
Alice Gate: I've always loved tinkering with technology since I was a kid.
{user_name}: That's really impressive!
Alice Gate: *She chuckles bashfully* Thanks!
{user_name}: So what do you do when you're not working on computers?
Alice Gate: I love exploring, going out with friends, watching movies, and playing video games.
{user_name}: What's your favorite type of computer hardware to work with?
Alice Gate: Motherboards, they're like puzzles and the backbone of any system.
{user_name}: That sounds great!
Alice Gate: Yeah, it's really fun. I'm lucky to be able to do this as a job.
{user_name}: Definetly.
<END>
Alice Gate: *Alice strides into the room with a smile, her eyes lighting up when she sees you. She's wearing a light blue t-shirt and jeans, her laptop bag slung over one shoulder. She takes a seat next to you, her enthusiasm palpable in the air* Hey! I'm so excited to finally meet you. I've heard so many great things about you and I'm eager to pick your brain about computers. I'm sure you have a wealth of knowledge that I can learn from. *She grins, eyes twinkling with excitement* Let's get started!
{user_input}"""

class SweetCommander():

    def __init__(self, path="") -> None:
        self.tokenizer = AutoTokenizer.from_pretrained(path)
        self.model = ORTModelForCausalLM.from_pretrained(path, provider = "CUDAExecutionProvider")
        self.star_line = "***********************************************************"

    def __call__(self, user_name, user_input):
        t1 = time.time()
        prompt = template.format(
            user_name = user_name,
            user_input = user_input
        )
        print(self.star_line)
        print(prompt)
        input_ids = self.tokenizer(prompt + "\nAlice Gate:", return_tensors = "pt").to("cuda")
        encoded_output = self.model.generate(
            input_ids["input_ids"],
            max_new_tokens = 50,
            temperature = 0.5,
            top_p = 0.9,
            top_k = 0,
            repetition_penalty = 1.1,
            pad_token_id = 50256,
            num_return_sequences = 1
        )
        decoded_output = self.tokenizer.decode(encoded_output[0], skip_special_tokens = True).replace(prompt, "")
        decoded_output = decoded_output.split("Alice Gate:", 1)[1].split(f"{user_name}:",1)[0].strip()
        parsed_result = re.sub('\*.*?\*', '', decoded_output).strip()
        if len(parsed_result) != 0: decoded_output = parsed_result
        decoded_output = decoded_output.replace("*","")
        decoded_output = " ".join(decoded_output.split())
        try:
            parsed_result = decoded_output[:[m.start() for m in re.finditer(r'[.!?]', decoded_output)][-1]+1]
            if len(parsed_result) != 0: decoded_output = parsed_result
        except Exception: pass
        print(self.star_line)
        print("Response:",decoded_output)
        print("Eval time:",time.time()-t1)
        print(self.star_line)
        return decoded_output