from transformers import AutoTokenizer, AutoModelForCausalLM
import torch
import time


class Gemma2B:
    def __init__(self):
        self.model_name = "google/gemma-2b-it"
        self.tokenizer = AutoTokenizer.from_pretrained(self.model_name)
        self.model = AutoModelForCausalLM.from_pretrained(self.model_name, torch_dtype=torch.bfloat16, )

    def inference_cpu(self, chat_template, max_new_tokens=200):
        chat = self.tokenizer.apply_chat_template(chat_template, tokenize=False,
                                                  add_generation_prompt=True)

        input_ids = self.tokenizer(chat, return_tensors="pt")
        outputs = self.model.generate(**input_ids, max_length=300, max_new_tokens=300)
        return self.tokenizer.decode(outputs[0])


if __name__ == "__main__":
    llm = Gemma2B()
    start_time_cpu = time.time()
    print(llm.inference_cpu(
        [
            {"role": "user", "content": f"hello"}]
    ))
    end_time_cpu = time.time()
    print(f"CPU Inference Time: {end_time_cpu - start_time_cpu}")