initial commit for the llm

app.py

@@ -0,0 +1,211 @@
+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
+
+# 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.
+"""
+
+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)]
+    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()
+    
+    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 = []
+        for chunk in chunks:
+            response = generate_response_for_chunk(chunk, history, temperature, max_new_tokens)
+            responses.append(response)
+        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}")

requirements.txt

@@ -0,0 +1,4 @@
+accelerate
+transformers
+SentencePiece
+bitsandbytes

utils.py

@@ -0,0 +1,160 @@
+import gradio as gr
+import os
+import torch
+from transformers import AutoConfig, AutoTokenizer, AutoModelForCausalLM, TextIteratorStreamer
+from threading import Thread
+from accelerate import init_empty_weights, infer_auto_device_map, disk_offload
+
+# Set environment variables
+HF_TOKEN = os.getenv("HF_TOKEN")
+
+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;
+}
+"""
+
+def initialize_model(model_name, max_memory=None):
+    device = torch.device('cpu')
+    
+    # Load model configuration
+    config = AutoConfig.from_pretrained(model_name)
+    
+    with init_empty_weights():
+        # Initialize model with empty weights
+        model = AutoModelForCausalLM.from_config(config)
+    
+    # Create device map based on memory constraints
+    device_map = infer_auto_device_map(
+        model, max_memory=max_memory, no_split_module_classes=["GPTNeoXLayer"], dtype="float16"
+    )
+    
+    # Determine if offloading is needed
+    needs_offloading = any(device == 'disk' for device in device_map.values())
+    
+    if needs_offloading:
+        # Load model for offloading
+        model = AutoModelForCausalLM.from_pretrained(
+            model_name, device_map=device_map, offload_folder="offload",
+            offload_state_dict=True, torch_dtype=torch.float16
+        )
+        offload_directory = "offload/"
+        # Offload model to disk
+        disk_offload(model=model, offload_dir=offload_directory)
+    else:
+        # Load model normally to specified device
+        model = AutoModelForCausalLM.from_pretrained(
+            model_name, torch_dtype=torch.float16
+        )
+        model.to(device)
+    
+    return model
+
+try:
+    # Initialize the model and tokenizer
+    model_name = "meta-llama/Meta-Llama-3-8B-Instruct"
+    model = initialize_model(model_name, max_memory={"cpu": "GiB"})
+    tokenizer = AutoTokenizer.from_pretrained(model_name, use_auth_token=HF_TOKEN)
+except Exception as e:
+    print(f"Error initializing model: {e}")
+    exit(1)
+
+terminators = [
+    tokenizer.eos_token_id,
+    tokenizer.convert_tokens_to_ids("")
+]
+
+def chat_llama3_8b(message: str, history: list, temperature: float, max_new_tokens: int) -> str:
+    """
+    Generate a streaming response using the llama3-8b model.
+    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.
+    """
+    conversation = []
+    message += " Extract all relevant keywords and add quantity from the following text and format the result in nested JSON:"
+    for user, assistant in history:
+        conversation.extend([{"role": "user", "content": user}, {"role": "assistant", "content": assistant}])
+    conversation.append({"role": "user", "content": message})
+
+    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,
+    )
+    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)
+        yield "".join(outputs)
+
+# 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=9012,
+                step=1,
+                value=512,
+                label="Max new tokens",
+                render=False
+            ),
+        ]
+    )
+    gr.Markdown(LICENSE)
+
+if __name__ == "__main__":
+    demo.launch(server_port=8000, share=True)