Update app.py

app.py

@@ -1,3 +1,76 @@
 import gradio as gr
+import spaces
+import torch
+from transformers import AutoModelForCausalLM, AutoTokenizer, pipeline
 
-gr.load("models/bigcode/starcoder2-15b").launch()
+# Model configuration
+CHECKPOINT = "bigcode/starcoder2-15b"
+DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
+
+# Load tokenizer and model (using bfloat16 for efficiency)
+@spaces.GPU(duration=120)  # Set duration to 120s to handle model loading/generation
+def load_model_and_generate(prompt, max_length=256, temperature=0.2, top_p=0.95):
+    try:
+        # Initialize tokenizer
+        tokenizer = AutoTokenizer.from_pretrained(CHECKPOINT)
+        
+        # Initialize model with bfloat16 for lower memory usage
+        model = AutoModelForCausalLM.from_pretrained(
+            CHECKPOINT,
+            torch_dtype=torch.bfloat16,
+            device_map="auto"
+        )
+        
+        # Create text generation pipeline
+        pipe = pipeline(
+            "text-generation",
+            model=model,
+            tokenizer=tokenizer,
+            device_map="auto",
+            torch_dtype=torch.bfloat16
+        )
+        
+        # Generate response
+        result = pipe(
+            prompt,
+            max_length=max_length,
+            temperature=temperature,
+            top_p=top_p,
+            num_return_sequences=1,
+            do_sample=True,
+            eos_token_id=tokenizer.eos_token_id,
+            pad_token_id=tokenizer.eos_token_id,
+            truncation=True
+        )
+        
+        generated_text = result[0]["generated_text"]
+        return generated_text
+    except Exception as e:
+        return f"Error: {str(e)}"
+
+# Gradio interface setup
+with gr.Blocks() as demo:
+    gr.Markdown("# StarCoder2-15B Code Generation")
+    gr.Markdown("Enter a code prompt (e.g., 'def print_hello_world():') to generate code using bigcode/starcoder2-15b.")
+    
+    # Input components
+    prompt = gr.Textbox(label="Code Prompt", placeholder="Enter your code prompt here...")
+    max_length = gr.Slider(50, 512, value=256, label="Max Length", step=1)
+    temperature = gr.Slider(0.1, 1.0, value=0.2, label="Temperature", step=0.1)
+    top_p = gr.Slider(0.1, 1.0, value=0.95, label="Top P", step=0.05)
+    
+    # Output component
+    output = gr.Textbox(label="Generated Code")
+    
+    # Submit button
+    submit_btn = gr.Button("Generate")
+    
+    # Connect button to function
+    submit_btn.click(
+        fn=load_model_and_generate,
+        inputs=[prompt, max_length, temperature, top_p],
+        outputs=output
+    )
+
+# Launch the Gradio app
+demo.launch()