wip

backend-developers-ltd · Sep 5, 2024 · f8d2ba9 · f8d2ba9
1 parent 963785d
commit f8d2ba9
Showing 1 changed file with 3 additions and 2 deletions.
diff --git a/src/compute_horde_prompt_gen/model.py b/src/compute_horde_prompt_gen/model.py
@@ -80,8 +80,8 @@ def tokenize(prompt: str) -> str:
     def tokenize_phi3(self, prompts: list[str], role: str) -> str:
         inputs = [{"role": role, "content": prompt} for prompt in prompts]
         inputs = self.tokenizer.apply_chat_template(
-            **inputs, add_generation_prompt=True, return_tensors="pt"
-        )
+            inputs, add_generation_prompt=True, return_tensors="pt"
+        ).to("cuda")
         return inputs
 
     def generate(
@@ -99,6 +99,7 @@ def generate(
             inputs = self.tokenize_phi3(prompts, role)
         else:
             raise ValueError(f"Unknown model {self.model_name}")
+        print(inputs)
 
         return self.model.generate(
             **inputs,