forked from nomic-ai/gpt4all
-
Notifications
You must be signed in to change notification settings - Fork 0
/
generate.py
55 lines (36 loc) · 1.83 KB
/
generate.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
from transformers import AutoModelForCausalLM, AutoTokenizer
from peft import PeftModelForCausalLM
from read import read_config
from argparse import ArgumentParser
import torch
import time
def generate(tokenizer, prompt, model, config):
input_ids = tokenizer(prompt, return_tensors="pt").input_ids.to(model.device)
outputs = model.generate(input_ids=input_ids, max_new_tokens=config["max_new_tokens"], temperature=config["temperature"])
decoded = tokenizer.decode(outputs[0], skip_special_tokens=True).strip()
return decoded[len(prompt):]
def setup_model(config):
model = AutoModelForCausalLM.from_pretrained(config["model_name"], device_map="auto", torch_dtype=torch.float16)
tokenizer = AutoTokenizer.from_pretrained(config["tokenizer_name"])
tokenizer.add_special_tokens({"bos_token": "<s>", "eos_token": "</s>"})
if config["lora"]:
model = PeftModelForCausalLM.from_pretrained(model, config["lora_path"], device_map="auto", torch_dtype=torch.float16)
model.to(dtype=torch.float16)
print(f"Mem needed: {model.get_memory_footprint() / 1024 / 1024 / 1024:.2f} GB")
return model, tokenizer
if __name__ == "__main__":
parser = ArgumentParser()
parser.add_argument("--config", type=str, required=True)
parser.add_argument("--prompt", type=str)
args = parser.parse_args()
config = read_config(args.config)
if config["prompt"] is None and args.prompt is None:
raise ValueError("Prompt is required either in config or as argument")
prompt = config["prompt"] if args.prompt is None else args.prompt
print("Setting up model")
model, tokenizer = setup_model(config)
print("Generating")
start = time.time()
generation = generate(tokenizer, prompt, model, config)
print(f"Done in {time.time() - start:.2f}s")
print(generation)