Views
No views yet
1import transformers
2import torch
3
4
5def fmt_prompt(prompt: str) -> str:
6 return f"""[Instructions]:\n{prompt}\n\n[Response]:"""
7
8
9if __name__ == "__main__":
10 model_name = "abacaj/mistral-7b-sft"
11 tokenizer = transformers.AutoTokenizer.from_pretrained(model_name)
12
13 model = (
14 transformers.AutoModelForCausalLM.from_pretrained(
15 model_name,
16 )
17 .to("cuda:0")
18 .eval()
19 )
20
21 prompt = "If A is greater than B and B is greater than C does that make A greater than C?"
22 prompt_input = fmt_prompt(prompt)
23 inputs = tokenizer(prompt_input, return_tensors="pt").to(model.device)
24 input_ids_cutoff = inputs.input_ids.size(dim=1)
25
26 with torch.no_grad():
27 generated_ids = model.generate(
28 **inputs,
29 use_cache=True,
30 max_new_tokens=512,
31 temperature=0.2,
32 top_p=0.95,
33 do_sample=True,
34 eos_token_id=tokenizer.eos_token_id,
35 pad_token_id=tokenizer.pad_token_id,
36 )
37
38 completion = tokenizer.decode(
39 generated_ids[0][input_ids_cutoff:],
40 skip_special_tokens=True,
41 )
42
43 print(completion)
