contenteaseAI sarjunshankar commited on
Commit
9ecba39
1 Parent(s): 0fe516c

llm-pr-branch (#1)

Browse files

- added code for llama_3.1 chat (b0984526920f5f4edcb6f16a483a9439ae209bb9)


Co-authored-by: Arjun Shankar <sarjunshankar@users.noreply.huggingface.co>

Files changed (2) hide show
  1. app.py +269 -47
  2. utils.py +160 -0
app.py CHANGED
@@ -1,63 +1,285 @@
1
  import gradio as gr
2
- from huggingface_hub import InferenceClient
 
 
 
 
 
 
 
3
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
4
  """
5
- For more information on `huggingface_hub` Inference API support, please check the docs: https://huggingface.co/docs/huggingface_hub/v0.22.2/en/guides/inference
 
 
 
 
 
6
  """
7
- client = InferenceClient("HuggingFaceH4/zephyr-7b-beta")
8
 
 
 
 
 
 
 
9
 
10
- def respond(
11
- message,
12
- history: list[tuple[str, str]],
13
- system_message,
14
- max_tokens,
15
- temperature,
16
- top_p,
17
- ):
18
- messages = [{"role": "system", "content": system_message}]
19
 
20
- for val in history:
21
- if val[0]:
22
- messages.append({"role": "user", "content": val[0]})
23
- if val[1]:
24
- messages.append({"role": "assistant", "content": val[1]})
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
25
 
26
- messages.append({"role": "user", "content": message})
 
 
 
 
27
 
28
- response = ""
 
 
 
29
 
30
- for message in client.chat_completion(
31
- messages,
32
- max_tokens=max_tokens,
33
- stream=True,
34
- temperature=temperature,
35
- top_p=top_p,
36
- ):
37
- token = message.choices[0].delta.content
38
 
39
- response += token
40
- yield response
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
41
 
42
  """
43
- For information on how to customize the ChatInterface, peruse the gradio docs: https://www.gradio.app/docs/chatinterface
44
- """
45
- demo = gr.ChatInterface(
46
- respond,
47
- additional_inputs=[
48
- gr.Textbox(value="You are a friendly Chatbot.", label="System message"),
49
- gr.Slider(minimum=1, maximum=2048, value=512, step=1, label="Max new tokens"),
50
- gr.Slider(minimum=0.1, maximum=4.0, value=0.7, step=0.1, label="Temperature"),
51
- gr.Slider(
52
- minimum=0.1,
53
- maximum=1.0,
54
- value=0.95,
55
- step=0.05,
56
- label="Top-p (nucleus sampling)",
57
- ),
58
- ],
59
- )
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
60
 
 
 
61
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
62
  if __name__ == "__main__":
63
- demo.launch()
 
 
 
 
1
  import gradio as gr
2
+ import os
3
+ import time
4
+ from transformers import AutoModelForCausalLM, AutoTokenizer, TextIteratorStreamer, BitsAndBytesConfig
5
+ import torch
6
+ from threading import Thread
7
+ import logging
8
+ import spaces
9
+ from functools import lru_cache
10
 
11
+ print(f"Is CUDA available: {torch.cuda.is_available()}")
12
+ # True
13
+ print(f"CUDA device: {torch.cuda.get_device_name(torch.cuda.current_device())}")
14
+
15
+ # Set up logging
16
+ logging.basicConfig(level=logging.INFO)
17
+ logger = logging.getLogger(__name__)
18
+
19
+ # Set an environment variable
20
+ HF_TOKEN = os.environ.get("HF_TOKEN", None)
21
+
22
+ DESCRIPTION = '''
23
+ <div>
24
+ <h1 style="text-align: center;">ContenteaseAI custom trained model</h1>
25
+ </div>
26
+ '''
27
+
28
+ LICENSE = """
29
+ <p/>
30
+ ---
31
+ For more information, visit our [website](https://contentease.ai).
32
  """
33
+
34
+ PLACEHOLDER = """
35
+ <div style="padding: 30px; text-align: center; display: flex; flex-direction: column; align-items: center;">
36
+ <h1 style="font-size: 28px; margin-bottom: 2px; opacity: 0.55;">ContenteaseAI Custom AI trained model</h1>
37
+ <p style="font-size: 18px; margin-bottom: 2px; opacity: 0.65;">Enter the text extracted from the PDF:</p>
38
+ </div>
39
  """
 
40
 
41
+ css = """
42
+ h1 {
43
+ text-align: center;
44
+ display: block;
45
+ }
46
+ """
47
 
48
+ # Load the tokenizer and model with quantization
49
+ model_id = "meta-llama/Meta-Llama-3.1-8B-Instruct"
50
+ bnb_config = BitsAndBytesConfig(
51
+ load_in_4bit=True,
52
+ bnb_4bit_use_double_quant=True,
53
+ bnb_4bit_quant_type="nf4",
54
+ bnb_4bit_compute_dtype=torch.bfloat16
55
+ )
 
56
 
57
+ @lru_cache(maxsize=1)
58
+ def load_model_and_tokenizer():
59
+ try:
60
+ start_time = time.time()
61
+ logger.info("Loading tokenizer...")
62
+ tokenizer = AutoTokenizer.from_pretrained(model_id)
63
+ logger.info("Loading model...")
64
+ model = AutoModelForCausalLM.from_pretrained(
65
+ model_id,
66
+ device_map="auto",
67
+ quantization_config=bnb_config,
68
+ torch_dtype=torch.bfloat16
69
+ )
70
+ model.generation_config.pad_token_id = tokenizer.pad_token_id
71
+ end_time = time.time()
72
+ logger.info(f"Model and tokenizer loaded successfully in {end_time - start_time} seconds.")
73
+ return model, tokenizer
74
+ except Exception as e:
75
+ logger.error(f"Error loading model or tokenizer: {e}")
76
+ raise
77
 
78
+ try:
79
+ model, tokenizer = load_model_and_tokenizer()
80
+ except Exception as e:
81
+ logger.error(f"Failed to load model and tokenizer: {e}")
82
+ raise
83
 
84
+ terminators = [
85
+ tokenizer.eos_token_id,
86
+ tokenizer.convert_tokens_to_ids("<|eot_id|>")
87
+ ]
88
 
89
+ # SYS_PROMPT = """
90
+ # Extract all relevant keywords and add quantity from the following text and format the result in nested JSON, ignoring personal details and focusing only on the scope of work as shown in the example:
91
+ # Good JSON example: {'lobby': {'frcm': {'replace': {'carpet': 1, 'carpet_pad': 1, 'base': 1, 'window_treatments': 1, 'artwork_and_decorative_accessories': 1, 'portable_lighting': 1, 'upholstered_furniture_and_decorative_pillows': 1, 'millwork': 1} } } }
92
+ # Bad JSON example: {'lobby': { 'frcm': { 'replace': [ 'carpet', 'carpet_pad', 'base', 'window_treatments', 'artwork_and_decorative_accessories', 'portable_lighting', 'upholstered_furniture_and_decorative_pillows', 'millwork'] } } }
93
+ # Make sure to fetch details from the provided text and ignore unnecessary information. The response should be in JSON format only, without any additional comments.
94
+ # """
 
 
95
 
96
+ SYS_PROMPT = """
97
+ Extract all relevant keywords and add quantities from the following text and format the result in nested JSON, ignoring personal details and focusing only on the area and furniture items as shown in the example. Each item should have a count, which will be set to 1 for simplicity. The response should be in JSON format only, without any additional comments.
98
+
99
+ Good JSON example:{
100
+ "Lobby Area/Entrance": {
101
+ "Vinyl wall covering": 1,
102
+ "Decorative hardwired lighting": 1
103
+ },
104
+ "Lobby": {
105
+ "Carpet, carpet pad, and base": 1,
106
+ "Window treatments": 1,
107
+ "Artwork and decorative accessories": 1,
108
+ "Portable lighting": 1,
109
+ "Upholstered furniture and decorative pillows": 1,
110
+ "Millwork": 1
111
+ }
112
+ }
113
+ Make sure to fetch details from the provided text and ignore unnecessary information. The response should be in JSON format only, without any additional comments.
114
+
115
+ Task:
116
+ Convert the provided extracted text into the JSON format described above.
117
+
118
+ Provided Text:
119
+
120
+ PROPERTY IMPROVEMENT PLAN
121
+ PREPARED FOR:
122
+ Springfield, IL
123
+ To be relicensed as Hilton Garden Inn
124
+ ...
125
+ Patios/The Terrace - Install patio decorative lighting. Install patio furniture. (lounge chairs, chaise, dining tables/chairs)
126
+ ...
127
+ Lobby Area - Replace carpet, carpet pad, and base. Replace window treatments. Replace artwork and decorative accessories. Replace portable lighting. (floor lamps, table lamps) Replace upholstered furniture and decorative pillows. Replace millwork. Replace the television(s).
128
+ ...
129
+ Registration Area - Replace vinyl wall covering. Replace hard surface floor covering. Replace artwork. Install new signature graphics on the back wall.
130
+ ...
131
+
132
+ Expected Output (JSON format):
133
+ {
134
+ "Patios/The Terrace": {
135
+ "Patio decorative lighting": 1,
136
+ "Lounge chairs": 1,
137
+ "Chaise": 1,
138
+ "Dining tables": 1,
139
+ "Dining chairs": 1,
140
+ "Patio furniture": 1
141
+ },
142
+ "Lobby Area": {
143
+ "Carpet, carpet pad, and base": 1,
144
+ "Window treatments": 1,
145
+ "Artwork and decorative accessories": 1,
146
+ "Portable lighting (floor lamps, table lamps)": 1,
147
+ "Upholstered furniture and decorative pillows": 1,
148
+ "Millwork": 1,
149
+ "Television(s)": 1
150
+ },
151
+ "Registration Area": {
152
+ "Vinyl wall covering": 1,
153
+ "Hard surface floor covering": 1,
154
+ "Artwork (new signature graphics on the back wall)": 1
155
+ }
156
+ }
157
 
158
  """
159
+ def chunk_text(text, chunk_size=5000):
160
+ """
161
+ Splits the input text into chunks of specified size.
162
+
163
+ Args:
164
+ text (str): The input text to be chunked.
165
+ chunk_size (int): The size of each chunk in tokens.
166
+
167
+ Returns:
168
+ list: A list of text chunks.
169
+ """
170
+ words = text.split()
171
+ chunks = [' '.join(words[i:i + chunk_size]) for i in range(0, len(words), chunk_size)]
172
+ logger.info(f"Total chunks created: {len(chunks)}")
173
+ return chunks
174
+
175
+ def combine_responses(responses):
176
+ """
177
+ Combines the responses from all chunks into a final output string.
178
+
179
+ Args:
180
+ responses (list): A list of responses from each chunk.
181
+
182
+ Returns:
183
+ str: The combined output string.
184
+ """
185
+ combined_output = " ".join(responses)
186
+ return combined_output
187
+
188
+ def generate_response_for_chunk(chunk, history, temperature, max_new_tokens):
189
+ start_time = time.time()
190
+ if len(history) == 0:
191
+ pass
192
+ else:
193
+ history.pop()
194
+ conversation = [{"role": "system", "content": SYS_PROMPT}]
195
+ for user, assistant in history:
196
+ conversation.extend([{"role": "user", "content": user}, {"role": "assistant", "content": assistant}])
197
+ conversation.append({"role": "user", "content": chunk})
198
+
199
+ input_ids = tokenizer.apply_chat_template(conversation, return_tensors="pt").to(model.device)
200
+
201
+ streamer = TextIteratorStreamer(tokenizer, timeout=10.0, skip_prompt=True, skip_special_tokens=True)
202
+
203
+ generate_kwargs = dict(
204
+ input_ids=input_ids,
205
+ streamer=streamer,
206
+ max_new_tokens=max_new_tokens,
207
+ do_sample=True,
208
+ temperature=temperature,
209
+ eos_token_id=terminators,
210
+ pad_token_id=tokenizer.eos_token_id
211
+ )
212
+ if temperature == 0:
213
+ generate_kwargs['do_sample'] = False
214
+
215
+ t = Thread(target=model.generate, kwargs=generate_kwargs)
216
+ t.start()
217
+
218
+ outputs = []
219
+ for text in streamer:
220
+ outputs.append(text)
221
+
222
+ end_time = time.time()
223
+ logger.info(f"Time taken for generating response for a chunk: {end_time - start_time} seconds")
224
+
225
+ return "".join(outputs)
226
+
227
+ @spaces.GPU(duration=110)
228
+ def chat_llama3_8b(message: str, history: list, temperature: float, max_new_tokens: int):
229
+ """
230
+ Generate a streaming response using the llama3-8b model with chunking.
231
+
232
+ Args:
233
+ message (str): The input message.
234
+ history (list): The conversation history used by ChatInterface.
235
+ temperature (float): The temperature for generating the response.
236
+ max_new_tokens (int): The maximum number of new tokens to generate.
237
+
238
+ Returns:
239
+ str: The generated response.
240
+ """
241
+ try:
242
+ start_time = time.time()
243
+
244
+ chunks = chunk_text(message)
245
+ responses = []
246
+ count=0
247
+ for chunk in chunks:
248
+ logger.info(f"Processing chunk {count+1}/{len(chunks)}")
249
+ response = generate_response_for_chunk(chunk, history, temperature, max_new_tokens)
250
+ responses.append(response)
251
+ count+=1
252
+ final_output = combine_responses(responses)
253
+
254
+ end_time = time.time()
255
+ logger.info(f"Total time taken for generating response: {end_time - start_time} seconds")
256
+
257
+ yield final_output
258
+ except Exception as e:
259
+ logger.error(f"Error generating response: {e}")
260
+ yield "An error occurred while generating the response. Please try again."
261
 
262
+ # Gradio block
263
+ chatbot = gr.Chatbot(height=450, placeholder=PLACEHOLDER, label='Gradio ChatInterface')
264
 
265
+ with gr.Blocks(fill_height=True, css=css) as demo:
266
+ gr.Markdown(DESCRIPTION)
267
+
268
+ gr.ChatInterface(
269
+ fn=chat_llama3_8b,
270
+ chatbot=chatbot,
271
+ fill_height=True,
272
+ additional_inputs_accordion=gr.Accordion(label="⚙️ Parameters", open=False, render=False),
273
+ additional_inputs=[
274
+ gr.Slider(minimum=0, maximum=1, step=0.1, value=0.95, label="Temperature", render=False),
275
+ gr.Slider(minimum=128, maximum=2000, step=1, value=700, label="Max new tokens", render=False),
276
+ ]
277
+ )
278
+
279
+ gr.Markdown(LICENSE)
280
+
281
  if __name__ == "__main__":
282
+ try:
283
+ demo.launch(show_error=True)
284
+ except Exception as e:
285
+ logger.error(f"Error launching Gradio demo: {e}")
utils.py ADDED
@@ -0,0 +1,160 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import gradio as gr
2
+ import os
3
+ import torch
4
+ from transformers import AutoConfig, AutoTokenizer, AutoModelForCausalLM, TextIteratorStreamer
5
+ from threading import Thread
6
+ from accelerate import init_empty_weights, infer_auto_device_map, disk_offload
7
+
8
+ # Set environment variables
9
+ HF_TOKEN = os.getenv("HF_TOKEN")
10
+
11
+ DESCRIPTION = '''
12
+ <div>
13
+ <h1 style="text-align: center;">ContenteaseAI custom trained model</h1>
14
+ </div>
15
+ '''
16
+
17
+ LICENSE = """
18
+ <p/>
19
+
20
+ ---
21
+ For more information, visit our [website](https://contentease.ai).
22
+ """
23
+
24
+ PLACEHOLDER = """
25
+ <div style="padding: 30px; text-align: center; display: flex; flex-direction: column; align-items: center;">
26
+ <h1 style="font-size: 28px; margin-bottom: 2px; opacity: 0.55;">ContenteaseAI Custom AI trained model</h1>
27
+ <p style="font-size: 18px; margin-bottom: 2px; opacity: 0.65;">Enter the text extracted from the PDF:</p>
28
+ </div>
29
+ """
30
+
31
+ css = """
32
+ h1 {
33
+ text-align: center;
34
+ display: block;
35
+ }
36
+ """
37
+
38
+ def initialize_model(model_name, max_memory=None):
39
+ device = torch.device('cpu')
40
+
41
+ # Load model configuration
42
+ config = AutoConfig.from_pretrained(model_name)
43
+
44
+ with init_empty_weights():
45
+ # Initialize model with empty weights
46
+ model = AutoModelForCausalLM.from_config(config)
47
+
48
+ # Create device map based on memory constraints
49
+ device_map = infer_auto_device_map(
50
+ model, max_memory=max_memory, no_split_module_classes=["GPTNeoXLayer"], dtype="float16"
51
+ )
52
+
53
+ # Determine if offloading is needed
54
+ needs_offloading = any(device == 'disk' for device in device_map.values())
55
+
56
+ if needs_offloading:
57
+ # Load model for offloading
58
+ model = AutoModelForCausalLM.from_pretrained(
59
+ model_name, device_map=device_map, offload_folder="offload",
60
+ offload_state_dict=True, torch_dtype=torch.float16
61
+ )
62
+ offload_directory = "offload/"
63
+ # Offload model to disk
64
+ disk_offload(model=model, offload_dir=offload_directory)
65
+ else:
66
+ # Load model normally to specified device
67
+ model = AutoModelForCausalLM.from_pretrained(
68
+ model_name, torch_dtype=torch.float16
69
+ )
70
+ model.to(device)
71
+
72
+ return model
73
+
74
+ try:
75
+ # Initialize the model and tokenizer
76
+ model_name = "meta-llama/Meta-Llama-3.1-8B-Instruct"
77
+ model = initialize_model(model_name, max_memory={"cpu": "GiB"})
78
+ tokenizer = AutoTokenizer.from_pretrained(model_name, use_auth_token=HF_TOKEN)
79
+ except Exception as e:
80
+ print(f"Error initializing model: {e}")
81
+ exit(1)
82
+
83
+ terminators = [
84
+ tokenizer.eos_token_id,
85
+ tokenizer.convert_tokens_to_ids("")
86
+ ]
87
+
88
+ def chat_llama3_8b(message: str, history: list, temperature: float, max_new_tokens: int) -> str:
89
+ """
90
+ Generate a streaming response using the llama3-8b model.
91
+ Args:
92
+ message (str): The input message.
93
+ history (list): The conversation history used by ChatInterface.
94
+ temperature (float): The temperature for generating the response.
95
+ max_new_tokens (int): The maximum number of new tokens to generate.
96
+ Returns:
97
+ str: The generated response.
98
+ """
99
+ conversation = []
100
+ message += " Extract all relevant keywords and add quantity from the following text and format the result in nested JSON:"
101
+ for user, assistant in history:
102
+ conversation.extend([{"role": "user", "content": user}, {"role": "assistant", "content": assistant}])
103
+ conversation.append({"role": "user", "content": message})
104
+
105
+ input_ids = tokenizer.apply_chat_template(conversation, return_tensors="pt").to(model.device)
106
+
107
+ streamer = TextIteratorStreamer(tokenizer, timeout=10.0, skip_prompt=True, skip_special_tokens=True)
108
+
109
+ generate_kwargs = dict(
110
+ input_ids=input_ids,
111
+ streamer=streamer,
112
+ max_new_tokens=max_new_tokens,
113
+ do_sample=True,
114
+ temperature=temperature,
115
+ eos_token_id=terminators,
116
+ )
117
+ if temperature == 0:
118
+ generate_kwargs['do_sample'] = False
119
+
120
+ t = Thread(target=model.generate, kwargs=generate_kwargs)
121
+ t.start()
122
+
123
+ outputs = []
124
+ for text in streamer:
125
+ outputs.append(text)
126
+ yield "".join(outputs)
127
+
128
+ # Gradio block
129
+ chatbot = gr.Chatbot(height=450, placeholder=PLACEHOLDER, label='Gradio ChatInterface')
130
+
131
+ with gr.Blocks(fill_height=True, css=css) as demo:
132
+ gr.Markdown(DESCRIPTION)
133
+ gr.ChatInterface(
134
+ fn=chat_llama3_8b,
135
+ chatbot=chatbot,
136
+ fill_height=True,
137
+ additional_inputs_accordion=gr.Accordion(label="⚙️ Parameters", open=False, render=False),
138
+ additional_inputs=[
139
+ gr.Slider(
140
+ minimum=0,
141
+ maximum=1,
142
+ step=0.1,
143
+ value=0.95,
144
+ label="Temperature",
145
+ render=False
146
+ ),
147
+ gr.Slider(
148
+ minimum=128,
149
+ maximum=9012,
150
+ step=1,
151
+ value=512,
152
+ label="Max new tokens",
153
+ render=False
154
+ ),
155
+ ]
156
+ )
157
+ gr.Markdown(LICENSE)
158
+
159
+ if __name__ == "__main__":
160
+ demo.launch(server_port=8000, share=True)