Skip to content

Commit

Permalink
wip
Browse files Browse the repository at this point in the history
  • Loading branch information
andreea-popescu-reef committed Sep 5, 2024
1 parent 1876b31 commit 618250d
Showing 2 changed files with 9 additions and 12 deletions.
15 changes: 6 additions & 9 deletions src/compute_horde_prompt_gen/model.py
Original file line number Diff line number Diff line change
@@ -14,10 +14,8 @@ def __init__(self):
pass

def generate(self, prompts: list[str], num_return_sequences: int, **_kwargs):
return [1 for _ in range(len(prompts) * num_return_sequences)]

def decode(self, _output):
return f"COPY PASTE INPUT PROMPT {PROMPT_ENDING} Here is the list of prompts:\nHow are you?\nDescribe something\nCount to ten\n"
content = f"COPY PASTE INPUT PROMPT {PROMPT_ENDING} Here is the list of prompts:\nHow are you?\nDescribe something\nCount to ten\n"
return [content for _ in range(len(prompts) * num_return_sequences)]


class GenerativeModel:
@@ -100,18 +98,17 @@ def generate(
else:
raise ValueError(f"Unknown model {self.model_name}")

return self.model.generate(
output = self.model.generate(
inputs,
max_new_tokens=max_new_tokens,
temperature=temperature,
num_return_sequences=num_return_sequences,
do_sample=True, # use sampling-based decoding
)

def decode(self, output):
if self.model_name == LLAMA3:
return self.tokenizer.decode(output, skip_special_tokens=True)
return [self.tokenizer.decode(x, skip_special_tokens=True) for x in output]
elif self.model_name == PHI3:
return "".join(self.tokenizer.batch_decode(output))
# return "".join(self.tokenizer.batch_decode(output))
return self.tokenizer.batch_decode(output)
else:
raise ValueError(f"Unknown model {self.model_name}")
6 changes: 3 additions & 3 deletions src/compute_horde_prompt_gen/run.py
Original file line number Diff line number Diff line change
@@ -41,9 +41,9 @@ def generate_prompts(

new_prompts = []
for j, sequence in enumerate(sequences):
output = model.decode(sequence)
log.info(f"\n\n{i=} output={output}\n\n")
generated_prompts = parse_output(output)
# output = model.decode(sequence)
# log.info(f"\n\n{i=} output={output}\n\n")
generated_prompts = parse_output(sequence)
log.debug(f"{i=} sequence={j} {generated_prompts=} from {output=}")

log.info(f"{i=} sequence={j} generated {len(generated_prompts)} prompts")

0 comments on commit 618250d

Please sign in to comment.