Update app.py

app.py

@@ -1,69 +1,28 @@
 import gradio as gr
 from transformers import AutoModelForCausalLM, AutoTokenizer, pipeline, BitsAndBytesConfig
+from threading import Thread
 
-description = """# <p style="text-align: center; color: #292b47;"> 🏎️ <span style='color: #3264ff;'>DeciCoder-6B:</span> A Fast Code Generation Model💨 </p>
-<span style='color: #292b47;'>Welcome to <a href="https://huggingface.co/Deci/DeciCoder-6B" style="color: #3264ff;">DeciCoder</a>! 
-DeciCoder-6B was trained on the Python, Java, Javascript, Rust, C++, C, and C# subset of the Starcoder Training Dataset, and it's released under the Apache 2.0 license. This model is capable of code-completion and instruction following. It surpasses CodeGen 2.5 7B, CodeLlama 7B, abd StarCoder 7B in its supported languages on HumanEval, and leads by 3 points in Python over StarCoderBase 15.5B."""
+description = """
+<p style="text-align: center; font-size: 24px; color: #292b47;">
+    <strong>🚀 <span style='color: #3264ff;'>DeciCoder-6B: Bigger, Faster, Stronger </span></strong>
+</p>
+<span style='color: #292b47;'>Welcome to the <a href="https://huggingface.co/Deci/DeciCoder-6B" style="color: #3264ff;">DeciCoder-6B playground</a>! DeciCoder-6B was trained on the Python, Java, Javascript, Rust, C++, C, and C# subset of the Starcoder Training Dataset, and it's released under the Apache 2.0 license. This model is capable of code-completion and instruction following. It surpasses CodeGen 2.5 7B, CodeLlama 7B, abd StarCoder 7B in its supported languages on HumanEval, and leads by 3 points in Python over StarCoderBase 15.5B.</span>
+"""
 
-GENERATION_TITLE= "<p style='font-size: 24px; color: #292b47;'>💻 Your generated code:</p>"
 
-def instantiate_huggingface_model(
-    model_name,
-    quantization_config=None,
-    device_map="auto",
-    use_cache=True,
-    trust_remote_code=None,
-    pad_token=None,
-    padding_side="left"
-):
-    """
-    Instantiate a HuggingFace model with optional quantization using the BitsAndBytes library.
-
-    Parameters:
-    - model_name (str): The name of the model to load from HuggingFace's model hub.
-    - quantization_config (BitsAndBytesConfig, optional): Configuration for model quantization.
-      If None, defaults to a pre-defined quantization configuration for 4-bit quantization.
-    - device_map (str, optional): Device placement strategy for model layers ('auto' by default).
-    - use_cache (bool, optional): Whether to cache model outputs (False by default).
-    - trust_remote_code (bool, optional): Whether to trust remote code for custom layers (True by default).
-    - pad_token (str, optional): The pad token to be used by the tokenizer. If None, uses the EOS token.
-    - padding_side (str, optional): The side on which to pad the sequences ('left' by default).
-
-    Returns:
-    - model (PreTrainedModel): The instantiated model ready for inference or fine-tuning.
-    - tokenizer (PreTrainedTokenizer): The tokenizer associated with the model.
-
-    The function will throw an exception if model loading fails.
-    """
-
-    # If quantization_config is not provided, use the default configuration
-    if quantization_config is None:
-        quantization_config = BitsAndBytesConfig(
-            load_in_8bit=True,
-            low_cpu_mem_usage=True,
-        )
-
-    model = AutoModelForCausalLM.from_pretrained(
-        model_name,
-        quantization_config=quantization_config,
-        device_map=device_map,
-        use_cache=use_cache,
-        trust_remote_code=trust_remote_code
-    )
+checkpoint = "Deci-early-access/DeciCoder-6B"
 
-    tokenizer = AutoTokenizer.from_pretrained(model_name, 
-                                              trust_remote_code=trust_remote_code)
+model = AutoModelForCausalLM.from_pretrained(checkpoint,
+                                             trust_remote_code=True,
+                                             device_map="auto",
+                                             low_cpu_mem_usage=True,
+                                             load_in_4bit=True)
 
-    if pad_token is not None:
-        tokenizer.pad_token = pad_token
-    else:
-        tokenizer.pad_token = tokenizer.eos_token
-    
-    tokenizer.padding_side = padding_side
+tokenizer = AutoTokenizer.from_pretrained(checkpoint, trust_remote_code=True)
 
-    return model, tokenizer
+tokenizer.pad_token = tokenizer.eos_token
 
-model, tokenizer = instantiate_huggingface_model("Deci-early-access/DeciCoder-6B", trust_remote_code=True)
+tokenizer.padding_side = "left"
 
 pipe = pipeline("text-generation",
                 model=model,
@@ -73,22 +32,6 @@ pipe = pipeline("text-generation",
                 temperature=1e-3,
 )
 
-def post_processing(prompt: str, completion: str) -> str:
-    """
-    Post-processes the generated code completion with HTML styling.
-
-    Args:
-        prompt (str): The input code prompt.
-        completion (str): The generated code completion.
-
-    Returns:
-        str: The HTML-styled code with prompt and completion.
-    """
-    completion = "<span style='color: #ff5b86;'>" + completion + "</span>"
-    prompt = "<span style='color: #7484b7;'>" + prompt + "</span>"
-    code_html = f"<br><hr><br><pre style='font-size: 12px'><code>{prompt}{completion}</code></pre><br><hr>"
-    return GENERATION_TITLE + code_html
-
 def code_generation(prompt: str) -> str:
     """
     Generates code based on the given prompt. Handles both regular and FIM (Fill-In-Missing) generation.
@@ -98,25 +41,25 @@ def code_generation(prompt: str) -> str:
 
     Returns:
         str: The HTML-styled generated code.
-    """    
+    """
     completion = pipe(prompt)[0]['generated_text']
-    completion = completion[len(prompt):]
-    return post_processing(prompt, completion)
+    return completion
+
 
 demo = gr.Blocks(
     css=".gradio-container {background-color: #FAFBFF; color: #292b47}"
 )
 with demo:
+    gr.Markdown(value=description)
+    with gr.Row():
+        code = gr.Code(lines=10, language="python", label="👨🏽‍💻 Input", value="def nth_element_in_fibonnaci(element):\n    \"\"\"Returns the nth element of the Fibonnaci sequence.\"\"\"")
+        output = gr.Code(label="💻 Generated code")
     with gr.Row():
-        _, colum_2, _ = gr.Column(scale=1), gr.Column(scale=6), gr.Column(scale=1)
-        with colum_2:
-            gr.Markdown(value=description)
-            code = gr.Code(lines=5, language="python", label="Input code", value="def nth_element_in_fibonnaci(element):\n    \"\"\"Returns the nth element of the Fibonnaci sequence.\"\"\"")
-            run = gr.Button(value="👨🏽‍💻 Generate code", size='lg')
-            output = gr.HTML(label="💻 Your generated code")
-            
+        run = gr.Button(value="👨🏽‍💻 Generate code")
+        clear = gr.Button("🗑️ Clear")
 
+    clear.click(lambda: (None, None), None, [code, output], queue=False)
     event = run.click(code_generation, [code], output)
     gr.HTML(label="Keep in touch", value="<img src='https://huggingface.co/spaces/Deci/DeciCoder-Demo/resolve/main/deci-coder-banner.png' alt='Keep in touch' style='display: block; color: #292b47; margin: auto; max-width: 800px;'>")
 
-demo.launch()
+demo.launch(debug=True)