wip

src/compute_horde_prompt_gen/model.py

@@ -21,6 +21,8 @@ def generate(self, prompts: list[str], num_return_sequences: int, **_kwargs):
 
 class GenerativeModel:
     def __init__(self, model_path: str, quantize: bool = False):
+        self.input_prompt_ending = None
+
         import torch
         from transformers import (
             AutoTokenizer,
@@ -50,11 +52,44 @@ def __init__(self, model_path: str, quantize: bool = False):
         )
 
     def tokenize(self, prompts: list[str], role: str) -> str:
-        pass
+        # set default padding token
+        self.tokenizer.pad_token = self.tokenizer.eos_token
+
+        role_templates = {
+            "system": "<|begin_of_text|><|start_header_id|>system<|end_header_id|>\n{{{{ {} }}}}<|eot_id|>",
+            "user": "<|start_header_id|>user<|end_header_id|>\n{{{{ {} }}}}<|eot_id|>",
+            "assistant": "<|start_header_id|>assistant<|end_header_id|>\n{{{{ {} }}}}<|eot_id|>",
+            "end": "<|start_header_id|>assistant<|end_header_id|>",
+        }
+
+        def tokenize(prompt: str) -> str:
+            msgs = [
+                {"role": "system", "content": role},
+                {"role": "user", "content": prompt},
+            ]
+            full_prompt = io.StringIO()
+            for msg in msgs:
+                full_prompt.write(role_templates[msg["role"]].format(msg["content"]))
+            full_prompt.write(role_templates["end"])
+            return full_prompt.getvalue()
+
+        inputs = [tokenize(prompt) for prompt in prompts]
+        inputs = self.tokenizer(inputs, return_tensors="pt", padding=True).to("cuda")
+        return inputs
 
     def decode(self, output) -> list[str]:
         pass
 
+    def decode(self, output) -> list[str]:
+        print(f"\nraw_output: {output}\n")
+        return [
+            strip_input(
+                self.tokenizer.decode(x, skip_special_tokens=True),
+                self.input_prompt_ending,
+            )
+            for x in output
+        ]
+
     def generate(
         self,
         prompts: list[str],
@@ -67,7 +102,7 @@ def generate(
         inputs = self.tokenize(prompts, role)
 
         output = self.model.generate(
-            inputs,
+            **inputs,
             max_new_tokens=max_new_tokens,
             temperature=temperature,
             num_return_sequences=num_return_sequences,
@@ -78,58 +113,12 @@ def generate(
 
 
 class Phi3(GenerativeModel):
-    def decode(self, output) -> list[str]:
-        print(f"\nraw_output: {output}\n")
-        # return [
-        #     strip_input(x, "<|assistant|>") for x in self.tokenizer.batch_decode(output)
-        # ]
-        return [
-            strip_input(
-                self.tokenizer.decode(x, skip_special_tokens=True), " }}assistant"
-            )
-            for x in output
-        ]
-
-    def tokenize(self, prompts: list[str], role: str) -> str:
-        inputs = [{"role": "user", "content": prompt} for prompt in prompts]
-        print(f"\ninputs: {inputs}\n")
-        inputs = self.tokenizer.apply_chat_template(
-            inputs, add_generation_prompt=True, return_tensors="pt"
-        ).to("cuda")
-        return inputs
+    def __init__(self, model_path: str, quantize: bool = False):
+        super().__init__(model_path, quantize)
+        self.input_prompt_ending = "assistant<|end_header_id|>"
 
 
 class Llama3(GenerativeModel):
-    def decode(self, output) -> list[str]:
-        return [
-            strip_input(
-                self.tokenizer.decode(x, skip_special_tokens=True), " }}assistant"
-            )
-            for x in output
-        ]
-
-    def tokenize(self, prompts: list[str], role: str) -> str:
-        # set default padding token
-        self.tokenizer.pad_token = self.tokenizer.eos_token
-
-        role_templates = {
-            "system": "<|begin_of_text|><|start_header_id|>system<|end_header_id|>\n{{{{ {} }}}}<|eot_id|>",
-            "user": "<|start_header_id|>user<|end_header_id|>\n{{{{ {} }}}}<|eot_id|>",
-            "assistant": "<|start_header_id|>assistant<|end_header_id|>\n{{{{ {} }}}}<|eot_id|>",
-            "end": "<|start_header_id|>assistant<|end_header_id|>",
-        }
-
-        def tokenize(prompt: str) -> str:
-            msgs = [
-                {"role": "system", "content": role},
-                {"role": "user", "content": prompt},
-            ]
-            full_prompt = io.StringIO()
-            for msg in msgs:
-                full_prompt.write(role_templates[msg["role"]].format(msg["content"]))
-            full_prompt.write(role_templates["end"])
-            return full_prompt.getvalue()
-
-        inputs = [tokenize(prompt) for prompt in prompts]
-        inputs = self.tokenizer(inputs, return_tensors="pt", padding=True).to("cuda")
-        return inputs
+    def __init__(self, model_path: str, quantize: bool = False):
+        super().__init__(model_path, quantize)
+        self.input_prompt_ending = " }}assistant"