forked from databricks/databricks-ml-examples
-
Notifications
You must be signed in to change notification settings - Fork 0
/
Copy path01_load_inference.py
190 lines (141 loc) · 10.5 KB
/
01_load_inference.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
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
# Databricks notebook source
# MAGIC %md
# MAGIC # MPT-7B-instruct Inference on Databricks
# MAGIC [MPT-7B-instruct](https://huggingface.co/mosaicml/mpt-7b-instruct) is a 7B parameters model for short-form instruction following. It is built by fine tuning [MPT-7B](https://huggingface.co/mosaicml/mpt-7b) which is a decoder-style transformer model pretrained on 1T tokens of English text and code. The model eliminates the context length limits by replacing positional embeddings with Attention with Linear Biases (ALiBi).
# MAGIC
# MAGIC Environment for this notebook:
# MAGIC - Runtime: 13.1 GPU ML Runtime
# MAGIC - Instance: `g5.4xlarge` on AWS, `Standard_NV36ads_A10_v5` (or `Standard_NC4as_T4_v3` and set `USE_TRITON=False`) on Azure
# COMMAND ----------
# MAGIC %md
# MAGIC ## Install required packages
# COMMAND ----------
USE_TRITON = True # TODO: Please change this to False if not using T4 GPUs
# COMMAND ----------
# Skip this step if running on Databricks runtime 13.2 GPU and above.
!wget -O /local_disk0/tmp/libcusparse-dev-11-7_11.7.3.50-1_amd64.deb https://developer.download.nvidia.com/compute/cuda/repos/ubuntu2204/x86_64/libcusparse-dev-11-7_11.7.3.50-1_amd64.deb && \
dpkg -i /local_disk0/tmp/libcusparse-dev-11-7_11.7.3.50-1_amd64.deb && \
wget -O /local_disk0/tmp/libcublas-dev-11-7_11.10.1.25-1_amd64.deb https://developer.download.nvidia.com/compute/cuda/repos/ubuntu2204/x86_64/libcublas-dev-11-7_11.10.1.25-1_amd64.deb && \
dpkg -i /local_disk0/tmp/libcublas-dev-11-7_11.10.1.25-1_amd64.deb && \
wget -O /local_disk0/tmp/libcusolver-dev-11-7_11.4.0.1-1_amd64.deb https://developer.download.nvidia.com/compute/cuda/repos/ubuntu2204/x86_64/libcusolver-dev-11-7_11.4.0.1-1_amd64.deb && \
dpkg -i /local_disk0/tmp/libcusolver-dev-11-7_11.4.0.1-1_amd64.deb && \
wget -O /local_disk0/tmp/libcurand-dev-11-7_10.2.10.91-1_amd64.deb https://developer.download.nvidia.com/compute/cuda/repos/ubuntu2204/x86_64/libcurand-dev-11-7_10.2.10.91-1_amd64.deb && \
dpkg -i /local_disk0/tmp/libcurand-dev-11-7_10.2.10.91-1_amd64.deb
# COMMAND ----------
# MAGIC %pip install xformers==0.0.20 einops==0.6.1 flash-attn==v1.0.3.post0 triton-pre-mlir@git+https://github.com/vchiley/triton.git@triton_pre_mlir#subdirectory=python
# COMMAND ----------
# MAGIC %md
# MAGIC ## Inference
# MAGIC The below snippets are adapted from [the model card of mpt-7b-instruct](https://huggingface.co/mosaicml/mpt-7b-instruct). The example in the model card should also work on Databricks with the same environment.
# COMMAND ----------
import transformers
import torch
# it is suggested to pin the revision commit hash and not change it for reproducibility because the uploader might change the model afterwards; you can find the commmit history of mpt-7b-instruct in https://huggingface.co/mosaicml/mpt-7b-instruct/commits/main
name = "mosaicml/mpt-7b-instruct"
config = transformers.AutoConfig.from_pretrained(
name,
trust_remote_code=True
)
config.init_device = 'cuda'
if USE_TRITON:
config.attn_config['attn_impl'] = 'triton'
model = transformers.AutoModelForCausalLM.from_pretrained(
name,
config=config,
torch_dtype=torch.bfloat16,
trust_remote_code=True,
cache_dir="/local_disk0/.cache/huggingface/",
revision="bbe7a55d70215e16c00c1825805b81e4badb57d7"
)
tokenizer = transformers.AutoTokenizer.from_pretrained("EleutherAI/gpt-neox-20b", padding_side="left")
generator = transformers.pipeline("text-generation",
model=model,
config=config,
tokenizer=tokenizer,
torch_dtype=torch.bfloat16,
device=0)
# COMMAND ----------
def generate_text(prompt, **kwargs):
if "max_new_tokens" not in kwargs:
kwargs["max_new_tokens"] = 512
kwargs.update(
{
"pad_token_id": tokenizer.eos_token_id,
"eos_token_id": tokenizer.eos_token_id,
}
)
template = "Below is an instruction that describes a task. Write a response that appropriately completes the request.\n\n###Instruction\n{instruction}\n\n### Response\n"
if isinstance(prompt, str):
full_prompt = template.format(instruction=prompt)
generated_text = generator(full_prompt, **kwargs)
generated_text = generated_text[0]["generated_text"]
elif isinstance(prompt, list):
full_prompts = list(map(lambda promp: template.format(instruction=promp), prompt))
outputs = generator(full_prompts, **kwargs)
generated_text = [out[0]["generated_text"] for out in outputs]
return generated_text
# COMMAND ----------
question = "What is a large language model?"
print(generate_text(question))
# COMMAND ----------
generate_kwargs = {
'max_new_tokens': 100,
'temperature': 0.2,
'top_p': 0.9,
'top_k': 50,
'repetition_penalty': 1.0,
'no_repeat_ngram_size': 0,
'use_cache': True,
'do_sample': True,
'eos_token_id': tokenizer.eos_token_id,
'pad_token_id': tokenizer.eos_token_id,
}
questions = ["what is ML?", "Name 10 colors."]
print(generate_text(questions, **generate_kwargs))
# COMMAND ----------
# Check that the generation quality when the context is long
long_input = """Provide a concise summary of the below passage.
Hannah Arendt was one of the seminal political thinkers of the twentieth century. The power and originality of her thinking was evident in works such as The Origins of Totalitarianism, The Human Condition, On Revolution and The Life of the Mind. In these works and in numerous essays she grappled with the most crucial political events of her time, trying to grasp their meaning and historical import, and showing how they affected our categories of moral and political judgment. What was required, in her view, was a new framework that could enable us to come to terms with the twin horrors of the twentieth century, Nazism and Stalinism. She provided such framework in her book on totalitarianism, and went on to develop a new set of philosophical categories that could illuminate the human condition and provide a fresh perspective on the nature of political life.
Although some of her works now belong to the classics of the Western tradition of political thought, she has always remained difficult to classify. Her political philosophy cannot be characterized in terms of the traditional categories of conservatism, liberalism, and socialism. Nor can her thinking be assimilated to the recent revival of communitarian political thought, to be found, for example, in the writings of A. MacIntyre, M. Sandel, C. Taylor and M. Walzer. Her name has been invoked by a number of critics of the liberal tradition, on the grounds that she presented a vision of politics that stood in opposition some key liberal principles. There are many strands of Arendt’s thought that could justify such a claim, in particular, her critique of representative democracy, her stress on civic engagement and political deliberation, her separation of morality from politics, and her praise of the revolutionary tradition. However, it would be a mistake to view Arendt as an anti-liberal thinker. Arendt was in fact a stern defender of constitutionalism and the rule of law, an advocate of fundamental human rights (among which she included not only the right to life, liberty, and freedom of expression, but also the right to action and to opinion), and a critic of all forms of political community based on traditional ties and customs, as well as those based on religious, ethnic, or racial identity.
Arendt’s political thought cannot, in this sense, be identified either with the liberal tradition or with the claims advanced by a number of its critics. Arendt did not conceive of politics as a means for the satisfaction of individual preferences, nor as a way to integrate individuals around a shared conception of the good. Her conception of politics is based instead on the idea of active citizenship, that is, on the value and importance of civic engagement and collective deliberation about all matters affecting the political community. If there is a tradition of thought with which Arendt can be identified, it is the classical tradition of civic republicanism originating in Aristotle and embodied in the writings of Machiavelli, Montesquieu, Jefferson, and Tocqueville. According to this tradition politics finds its authentic expression whenever citizens gather together in a public space to deliberate and decide about matters of collective concern. Political activity is valued not because it may lead to agreement or to a shared conception of the good, but because it enables each citizen to exercise his or her powers of agency, to develop the capacities for judgment and to attain by concerted action some measure of political efficacy."""
def get_num_tokens(text):
inputs = tokenizer(text, return_tensors="pt").input_ids.to("cuda")
return inputs.shape[1]
print('number of tokens for input:', get_num_tokens(long_input))
results = generate_text([long_input], max_new_tokens=200)
print(results[0])
# COMMAND ----------
# MAGIC %md
# MAGIC ## Measure inference speed
# MAGIC Text generation speed is often measured with token/s, which is the average number of tokens that are generated by the model per second.
# COMMAND ----------
import time
import logging
import re
def get_gen_text_throughput(prompt, **kwargs):
if "max_new_tokens" not in kwargs:
kwargs["max_new_tokens"] = 512
kwargs.update(
{
"pad_token_id": tokenizer.eos_token_id,
"eos_token_id": tokenizer.eos_token_id,
"return_tensors": True, # make the pipeline return token ids instead of decoded text to get the number of generated tokens
}
)
template = "Below is an instruction that describes a task. Write a response that appropriately completes the request.\n\n###Instruction\n{instruction}\n\n### Response\n"
full_prompt = template.format(instruction=prompt)
# measure the time it takes for text generation
start = time.time()
generated_text = generator(full_prompt, **kwargs)
duration = time.time() - start
num_input_tokens = get_num_tokens(full_prompt)
# get the number of generated tokens
n_tokens = len(generated_text[0]["generated_token_ids"])
# show the generated text
generated_text = tokenizer.decode(generated_text[0]["generated_token_ids"],
skip_special_tokens=True)
return ((n_tokens - num_input_tokens) / duration, (n_tokens - num_input_tokens), generated_text)
# COMMAND ----------
throughput, n_tokens, result = get_gen_text_throughput("What is ML?", max_new_tokens=200)
print(f"{throughput} tokens/sec, {n_tokens} tokens (not including prompt)")
print(result)