StronLu commited on
Commit
1674ca8
·
1 Parent(s): 5ae5e10

Create model.py

Browse files
Files changed (1) hide show
  1. model.py +74 -0
model.py ADDED
@@ -0,0 +1,74 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ from threading import Thread
2
+ from typing import Iterator
3
+
4
+ import torch
5
+ from transformers import AutoConfig, AutoModelForCausalLM, AutoTokenizer, TextIteratorStreamer
6
+
7
+ model_id = 'meta-llama/Llama-2-13b-chat-hf'
8
+
9
+ if torch.cuda.is_available():
10
+ config = AutoConfig.from_pretrained(model_id)
11
+ config.pretraining_tp = 1
12
+ model = AutoModelForCausalLM.from_pretrained(
13
+ model_id,
14
+ config=config,
15
+ torch_dtype=torch.float16,
16
+ load_in_4bit=True,
17
+ device_map='auto'
18
+ )
19
+ else:
20
+ model = None
21
+ tokenizer = AutoTokenizer.from_pretrained(model_id)
22
+
23
+
24
+ def get_prompt(message: str, chat_history: list[tuple[str, str]],
25
+ system_prompt: str) -> str:
26
+ texts = [f'<s>[INST] <<SYS>>\n{system_prompt}\n<</SYS>>\n\n']
27
+ # The first user input is _not_ stripped
28
+ do_strip = False
29
+ for user_input, response in chat_history:
30
+ user_input = user_input.strip() if do_strip else user_input
31
+ do_strip = True
32
+ texts.append(f'{user_input} [/INST] {response.strip()} </s><s>[INST] ')
33
+ message = message.strip() if do_strip else message
34
+ texts.append(f'{message} [/INST]')
35
+ return ''.join(texts)
36
+
37
+
38
+ def get_input_token_length(message: str, chat_history: list[tuple[str, str]], system_prompt: str) -> int:
39
+ prompt = get_prompt(message, chat_history, system_prompt)
40
+ input_ids = tokenizer([prompt], return_tensors='np', add_special_tokens=False)['input_ids']
41
+ return input_ids.shape[-1]
42
+
43
+
44
+ def run(message: str,
45
+ chat_history: list[tuple[str, str]],
46
+ system_prompt: str,
47
+ max_new_tokens: int = 1024,
48
+ temperature: float = 0.8,
49
+ top_p: float = 0.95,
50
+ top_k: int = 50) -> Iterator[str]:
51
+ prompt = get_prompt(message, chat_history, system_prompt)
52
+ inputs = tokenizer([prompt], return_tensors='pt', add_special_tokens=False).to('cuda')
53
+
54
+ streamer = TextIteratorStreamer(tokenizer,
55
+ timeout=10.,
56
+ skip_prompt=True,
57
+ skip_special_tokens=True)
58
+ generate_kwargs = dict(
59
+ inputs,
60
+ streamer=streamer,
61
+ max_new_tokens=max_new_tokens,
62
+ do_sample=True,
63
+ top_p=top_p,
64
+ top_k=top_k,
65
+ temperature=temperature,
66
+ num_beams=1,
67
+ )
68
+ t = Thread(target=model.generate, kwargs=generate_kwargs)
69
+ t.start()
70
+
71
+ outputs = []
72
+ for text in streamer:
73
+ outputs.append(text)
74
+ yield ''.join(outputs)