Create llm.py

llm.py

@@ -0,0 +1,29 @@
+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}")