Views
No views yet
1from peft import PeftModel
2import torch
3from transformers import AutoTokenizer, AutoModelForCausalLM, BitsAndBytesConfig, GenerationConfig
4model_id = "EleutherAI/gpt-neox-20b"
5tokenizer = AutoTokenizer.from_pretrained(model_id)
6bnb_config = BitsAndBytesConfig(
7 load_in_4bit=True,
8 bnb_4bit_use_double_quant=True,
9 bnb_4bit_quant_type="nf4",
10 bnb_4bit_compute_dtype=torch.bfloat16
11)
12model = AutoModelForCausalLM.from_pretrained(model_id, quantization_config=bnb_config, device_map="auto")
13model = PeftModel.from_pretrained(model, "myzens/AlpaGo")
14
15#You can change Here.
16PROMPT = """Below is an instruction that describes a task. Write a response that appropriately completes the request.
17### Instruction:
18Write a short story about a lost key that unlocks a mysterious door.
19### Response:"""
20
21inputs = tokenizer(PROMPT, return_tensors="pt")
22input_ids = inputs["input_ids"].cuda()
23
24generation_config = GenerationConfig(
25 temperature=0.6,
26 top_p=0.95,
27 repetition_penalty=1.15,
28
29)
30
31print("Generating...")
32generation_output = model.generate(
33 input_ids=input_ids,
34 generation_config=generation_config,
35 return_dict_in_generate=True,
36 output_scores=True,
37 max_new_tokens=256,
38 eos_token_id=tokenizer.eos_token_id,
39 pad_token_id=tokenizer.pad_token_id,
40)
41
42for s in generation_output.sequences:
43 print(tokenizer.decode(s))
44