Update app.py

app.py

@@ -1,129 +1,79 @@
 import gradio as gr
-import torch
 from transformers import AutoModelForCausalLM, AutoTokenizer
+import torch
 
-# Load model with optimization for free GPU
 model_name = "mistralai/Devstral-Small-2505"
-tokenizer = AutoTokenizer.from_pretrained(model_name)
 
-# Optimization for free GPU
+tokenizer = AutoTokenizer.from_pretrained(model_name, use_fast=False)
 model = AutoModelForCausalLM.from_pretrained(
-    model_name,
-    torch_dtype=torch.float16,  # Use half precision
-    device_map="auto",          # Automatically distribute across available devices
-    load_in_8bit=True,          # 8-bit quantization to save memory
+    model_name, 
+    torch_dtype=torch.float16 if torch.cuda.is_available() else torch.float32
 )
+device = "cuda" if torch.cuda.is_available() else "cpu"
+model = model.to(device)
 
-def generate_text(prompt, max_length=512, temperature=0.7, top_p=0.9):
-    inputs = tokenizer(prompt, return_tensors="pt").to(model.device)
-    
+def code_completion(prompt, max_new_tokens=128, temperature=0.2):
+    if not prompt.strip():
+        return "Please enter some code to complete."
+    inputs = tokenizer(prompt, return_tensors="pt").to(device)
     with torch.no_grad():
         outputs = model.generate(
-            inputs.input_ids,
-            max_length=max_length,
+            **inputs,
+            max_new_tokens=max_new_tokens,
             temperature=temperature,
-            top_p=top_p,
             do_sample=True,
             pad_token_id=tokenizer.eos_token_id
         )
-    
-    response = tokenizer.decode(outputs[0], skip_special_tokens=True)
-    return response
+    generated = tokenizer.decode(outputs[0], skip_special_tokens=True)
+    return generated[len(prompt):]
 
-# Custom CSS for a more beautiful interface
 custom_css = """
-.container {
-    max-width: 900px;
-    margin: auto;
-    padding-top: 1.5rem;
-}
-.title {
+body {background: #f7f8fa;}
+.gradio-container {background: #f7f8fa;}
+h1, h2, h3, h4, h5, h6 {font-family: 'Inter', sans-serif;}
+#main-title {
     text-align: center;
-    color: #2a4365;
+    font-weight: 800;
+    font-size: 2.3em;
+    margin-bottom: 0.2em;
+    letter-spacing: -1px;
+    color: #222;
 }
-.subtitle {
+#subtitle {
     text-align: center;
-    color: #4a5568;
-    margin-bottom: 2rem;
-}
-footer {display: none !important;}
-.gradio-container {
-    font-family: 'Inter', -apple-system, BlinkMacSystemFont, sans-serif;
-}
-.gr-button {
-    border-radius: 8px !important;
-}
-.gr-button.gr-button-lg {
-    padding: 8px 18px;
-}
-.gr-input, .gr-textarea {
-    border-radius: 8px !important;
-}
-.gr-form {
-    border-radius: 12px !important;
-    padding: 20px !important;
-    background: #f8fafc !important;
-}
-.advanced-options {
-    margin-top: 1rem;
-    padding: 1rem;
-    border-radius: 8px;
-    background: #edf2f7;
+    color: #6c6f7a;
+    font-size: 1.1em;
+    margin-bottom: 2em;
 }
+.gr-box {border-radius: 16px;}
 """
 
-# Create enhanced minimalist interface
-with gr.Blocks(css=custom_css) as demo:
-    with gr.Column(elem_classes="container"):
-        gr.Markdown("# Devstral-Small-2505 Demo", elem_classes="title")
-        gr.Markdown("Experience the power of Mistral's latest code-focused model", elem_classes="subtitle")
-        
-        with gr.Box():
+with gr.Blocks(css=custom_css, theme=gr.themes.Soft()) as demo:
+    gr.Markdown(
+        """
+        <h1 id="main-title">Devstral Code Autocomplete</h1>
+        <div id="subtitle">Minimal, beautiful code completion powered by <b>Devstral</b></div>
+        """)
+    with gr.Row():
+        with gr.Column(scale=1):
             prompt = gr.Textbox(
-                placeholder="Enter your prompt here...", 
-                label="Input",
-                lines=5
+                label="Your code prompt",
+                lines=10,
+                placeholder="def quicksort(arr):\n    \"\"\"Sort the array using quicksort algorithm.\"\"\"\n    if len(arr) <= 1:\n        return arr\n    pivot = arr[len(arr) // 2]\n    ",
+                show_copy_button=True,
+                autofocus=True
             )
-            
             with gr.Row():
-                submit_btn = gr.Button("Generate", variant="primary", scale=2)
-                clear_btn = gr.Button("Clear", scale=1)
-            
-            with gr.Accordion("Advanced Options", open=False):
-                with gr.Row():
-                    with gr.Column():
-                        max_length = gr.Slider(
-                            minimum=64, maximum=2048, value=512, step=64,
-                            label="Maximum Length"
-                        )
-                    with gr.Column():
-                        temperature = gr.Slider(
-                            minimum=0.1, maximum=1.0, value=0.7, step=0.1,
-                            label="Temperature (Creativity)"
-                        )
-                with gr.Row():
-                    with gr.Column():
-                        top_p = gr.Slider(
-                            minimum=0.1, maximum=1.0, value=0.9, step=0.05,
-                            label="Top-p (Nucleus Sampling)"
-                        )
-            
-            output = gr.Textbox(
-                label="Model Output", 
+                max_tokens = gr.Slider(16, 256, value=128, step=8, label="Max new tokens")
+                temperature = gr.Slider(0.1, 1.0, value=0.2, step=0.05, label="Temperature")
+            btn = gr.Button("Generate Completion", elem_id="generate-btn")
+        with gr.Column(scale=1):
+            output = gr.Code(
+                label="Generated code",
+                language="python",
                 lines=12,
-                show_copy_button=True
-            )
-        
-        with gr.Row():
-            gr.Markdown(
-                "⚡ Powered by Mistral's Devstral-Small-2505 - Optimized for code generation and technical tasks"
+                interactive=False
             )
-    
-    submit_btn.click(
-        generate_text, 
-        inputs=[prompt, max_length, temperature, top_p], 
-        outputs=output
-    )
-    clear_btn.click(lambda: "", None, prompt)
+    btn.click(code_completion, inputs=[prompt, max_tokens, temperature], outputs=output)
 
 demo.launch()