Spaces:
Paused
Paused
File size: 9,976 Bytes
5f3e959 e1f5561 5f3e959 |
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 191 192 193 194 195 196 197 198 199 200 201 202 203 204 205 206 207 208 209 210 211 212 213 214 215 216 217 218 219 220 221 222 223 224 225 226 227 228 229 230 231 232 233 234 235 236 237 238 239 240 241 242 243 244 245 246 247 248 249 250 251 252 253 254 255 256 257 258 259 260 261 262 263 264 265 266 267 268 269 270 271 272 273 274 275 276 277 278 279 280 281 282 283 284 285 286 |
import gradio as gr
import os
import time
from transformers import AutoModelForCausalLM, AutoTokenizer, TextIteratorStreamer, BitsAndBytesConfig
import torch
from threading import Thread
import logging
import spaces
from functools import lru_cache
print(f"Is CUDA available: {torch.cuda.is_available()}")
# True
print(f"CUDA device: {torch.cuda.get_device_name(torch.cuda.current_device())}")
# Set up logging
logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)
# Set an environment variable
HF_TOKEN = os.environ.get("HF_TOKEN", None)
DESCRIPTION = '''
<div>
<h1 style="text-align: center;">ContenteaseAI custom trained model</h1>
</div>
'''
LICENSE = """
<p/>
---
For more information, visit our [website](https://contentease.ai).
"""
PLACEHOLDER = """
<div style="padding: 30px; text-align: center; display: flex; flex-direction: column; align-items: center;">
<h1 style="font-size: 28px; margin-bottom: 2px; opacity: 0.55;">ContenteaseAI Custom AI trained model</h1>
<p style="font-size: 18px; margin-bottom: 2px; opacity: 0.65;">Enter the text extracted from the PDF:</p>
</div>
"""
css = """
h1 {
text-align: center;
display: block;
}
"""
# Load the tokenizer and model with quantization
model_id = "meta-llama/Meta-Llama-3-8B-Instruct"
bnb_config = BitsAndBytesConfig(
load_in_4bit=True,
bnb_4bit_use_double_quant=True,
bnb_4bit_quant_type="nf4",
bnb_4bit_compute_dtype=torch.bfloat16
)
@lru_cache(maxsize=1)
def load_model_and_tokenizer():
try:
start_time = time.time()
logger.info("Loading tokenizer...")
tokenizer = AutoTokenizer.from_pretrained(model_id)
logger.info("Loading model...")
model = AutoModelForCausalLM.from_pretrained(
model_id,
device_map="auto",
quantization_config=bnb_config,
torch_dtype=torch.bfloat16
)
model.generation_config.pad_token_id = tokenizer.pad_token_id
end_time = time.time()
logger.info(f"Model and tokenizer loaded successfully in {end_time - start_time} seconds.")
return model, tokenizer
except Exception as e:
logger.error(f"Error loading model or tokenizer: {e}")
raise
try:
model, tokenizer = load_model_and_tokenizer()
except Exception as e:
logger.error(f"Failed to load model and tokenizer: {e}")
raise
terminators = [
tokenizer.eos_token_id,
tokenizer.convert_tokens_to_ids("<|eot_id|>")
]
# SYS_PROMPT = """
# 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:
# 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} } } }
# Bad JSON example: {'lobby': { 'frcm': { 'replace': [ 'carpet', 'carpet_pad', 'base', 'window_treatments', 'artwork_and_decorative_accessories', 'portable_lighting', 'upholstered_furniture_and_decorative_pillows', 'millwork'] } } }
# 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.
# """
SYS_PROMPT = """
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.
Good JSON example:{
"Lobby Area/Entrance": {
"Vinyl wall covering": 1,
"Decorative hardwired lighting": 1
},
"Lobby": {
"Carpet, carpet pad, and base": 1,
"Window treatments": 1,
"Artwork and decorative accessories": 1,
"Portable lighting": 1,
"Upholstered furniture and decorative pillows": 1,
"Millwork": 1
}
}
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.
Task:
Convert the provided extracted text into the JSON format described above.
Provided Text:
PROPERTY IMPROVEMENT PLAN
PREPARED FOR:
Springfield, IL
To be relicensed as Hilton Garden Inn
...
Patios/The Terrace - Install patio decorative lighting. Install patio furniture. (lounge chairs, chaise, dining tables/chairs)
...
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).
...
Registration Area - Replace vinyl wall covering. Replace hard surface floor covering. Replace artwork. Install new signature graphics on the back wall.
...
Expected Output (JSON format):
{
"Patios/The Terrace": {
"Patio decorative lighting": 1,
"Lounge chairs": 1,
"Chaise": 1,
"Dining tables": 1,
"Dining chairs": 1,
"Patio furniture": 1
},
"Lobby Area": {
"Carpet, carpet pad, and base": 1,
"Window treatments": 1,
"Artwork and decorative accessories": 1,
"Portable lighting (floor lamps, table lamps)": 1,
"Upholstered furniture and decorative pillows": 1,
"Millwork": 1,
"Television(s)": 1
},
"Registration Area": {
"Vinyl wall covering": 1,
"Hard surface floor covering": 1,
"Artwork (new signature graphics on the back wall)": 1
}
}
"""
def chunk_text(text, chunk_size=5000):
"""
Splits the input text into chunks of specified size.
Args:
text (str): The input text to be chunked.
chunk_size (int): The size of each chunk in tokens.
Returns:
list: A list of text chunks.
"""
words = text.split()
chunks = [' '.join(words[i:i + chunk_size]) for i in range(0, len(words), chunk_size)]
logger.info(f"Total chunks created: {len(chunks)}")
return chunks
def combine_responses(responses):
"""
Combines the responses from all chunks into a final output string.
Args:
responses (list): A list of responses from each chunk.
Returns:
str: The combined output string.
"""
combined_output = " ".join(responses)
return combined_output
def generate_response_for_chunk(chunk, history, temperature, max_new_tokens):
start_time = time.time()
if len(history) == 0:
pass
else:
history.pop()
conversation = [{"role": "system", "content": SYS_PROMPT}]
for user, assistant in history:
conversation.extend([{"role": "user", "content": user}, {"role": "assistant", "content": assistant}])
conversation.append({"role": "user", "content": chunk})
input_ids = tokenizer.apply_chat_template(conversation, return_tensors="pt").to(model.device)
streamer = TextIteratorStreamer(tokenizer, timeout=10.0, skip_prompt=True, skip_special_tokens=True)
generate_kwargs = dict(
input_ids=input_ids,
streamer=streamer,
max_new_tokens=max_new_tokens,
do_sample=True,
temperature=temperature,
eos_token_id=terminators,
pad_token_id=tokenizer.eos_token_id
)
if temperature == 0:
generate_kwargs['do_sample'] = False
t = Thread(target=model.generate, kwargs=generate_kwargs)
t.start()
outputs = []
for text in streamer:
outputs.append(text)
end_time = time.time()
logger.info(f"Time taken for generating response for a chunk: {end_time - start_time} seconds")
return "".join(outputs)
# @spaces.GPU(duration=110)
def chat_llama3_8b(message: str, history: list, temperature: float, max_new_tokens: int):
"""
Generate a streaming response using the llama3-8b model with chunking.
Args:
message (str): The input message.
history (list): The conversation history used by ChatInterface.
temperature (float): The temperature for generating the response.
max_new_tokens (int): The maximum number of new tokens to generate.
Returns:
str: The generated response.
"""
try:
start_time = time.time()
chunks = chunk_text(message)
responses = []
count=0
for chunk in chunks:
logger.info(f"Processing chunk {count+1}/{len(chunks)}")
response = generate_response_for_chunk(chunk, history, temperature, max_new_tokens)
responses.append(response)
count+=1
final_output = combine_responses(responses)
end_time = time.time()
logger.info(f"Total time taken for generating response: {end_time - start_time} seconds")
yield final_output
except Exception as e:
logger.error(f"Error generating response: {e}")
yield "An error occurred while generating the response. Please try again."
# Gradio block
chatbot = gr.Chatbot(height=450, placeholder=PLACEHOLDER, label='Gradio ChatInterface')
with gr.Blocks(fill_height=True, css=css) as demo:
gr.Markdown(DESCRIPTION)
gr.ChatInterface(
fn=chat_llama3_8b,
chatbot=chatbot,
fill_height=True,
additional_inputs_accordion=gr.Accordion(label="⚙️ Parameters", open=False, render=False),
additional_inputs=[
gr.Slider(minimum=0, maximum=1, step=0.1, value=0.95, label="Temperature", render=False),
gr.Slider(minimum=128, maximum=2000, step=1, value=700, label="Max new tokens", render=False),
]
)
gr.Markdown(LICENSE)
if __name__ == "__main__":
try:
demo.launch(show_error=True)
except Exception as e:
logger.error(f"Error launching Gradio demo: {e}")
|