rhysjones

/

phi-2-orange

like 42

Hey - there are a few recent updates to Phi-2 to base it on HF's transformer updates that post-date this model's Phi-2 version. One of the updates seems to fix the support for beam search. I'll look to update in the future with this and also allow it to run without needing trust_remote_code.

In the meantime, the other recent Phi-2 update is to make the eos token explicit. If you add this to your code, it should reliably finish on the EOS token:

import torch
import transformers
import time

model = transformers.AutoModelForCausalLM.from_pretrained(
    "rhysjones/phi-2-orange",
    trust_remote_code=True
)
tokenizer = transformers.AutoTokenizer.from_pretrained("rhysjones/phi-2-orange")


SYSTEM_PROMPT = "You are an AI assistant. You will be given a task. You must generate a short answer."
input_text = f"""<|im_start|>system
You are a helpful assistant that gives short answers.<|im_end|>
<|im_start|>user
Give me the first 3 prime numbers.<|im_end|>
<|im_start|>assistant
"""

generation_config = transformers.GenerationConfig(
    eos_token_id = 50256
)

t1 = time.time()
with torch.no_grad():
    outputs = model.generate(
        tokenizer(input_text, return_tensors="pt")['input_ids'],
        max_new_tokens=200,
        generation_config=generation_config,
        num_beams = 1
    )
print(tokenizer.decode(outputs[0], skip_special_tokens=True))
print(time.time()-t1)

Newer version also never stops, it is better to add a stoping criteria, for example for token <|im_end|>