wip

backend-developers-ltd · Sep 5, 2024 · 735a91a · 735a91a
1 parent 6b04b41
commit 735a91a
Show file tree

Hide file tree

Showing 2 changed files with 7 additions and 2 deletions.
diff --git a/src/compute_horde_prompt_gen/model.py b/src/compute_horde_prompt_gen/model.py
@@ -109,4 +109,9 @@ def generate(
         )
 
     def decode(self, output):
-        return self.tokenizer.decode(output, skip_special_tokens=True)
+        if self.model_name == LLAMA3:
+            return self.tokenizer.decode(output, skip_special_tokens=True)
+        elif self.model_name == PHI3:
+            return self.tokenizer.batch_decode(output)
+        else:
+            raise ValueError(f"Unknown model {self.model_name}")
diff --git a/src/compute_horde_prompt_gen/run.py b/src/compute_horde_prompt_gen/run.py
@@ -42,7 +42,7 @@ def generate_prompts(
         new_prompts = []
         for j, sequence in enumerate(sequences):
             output = model.decode(sequence)
-            log.info(f"{i=} output={output}")
+            log.info(f"\n\n{i=} output={output}\n\n")
             generated_prompts = parse_output(output)
             log.debug(f"{i=} sequence={j} {generated_prompts=} from {output=}")