Update app.py

app.py

@@ -1,59 +1,81 @@
-
-import gradio as gr
+import os
 import torch
-from transformers import AutoModelForCausalLM, AutoTokenizer
+import gradio as gr
+from transformers import AutoModelForCausalLM, AutoTokenizer, BitsAndBytesConfig
+
+# --- Config ---
+MODEL_PATH = os.getenv("MODEL_PATH", "WhiteRabbitNeo/WhiteRabbitNeo-2.5-Qwen-2.5-Coder-7B")
+LOAD_IN_4BIT = os.getenv("LOAD_IN_4BIT", "true").lower() == "true"
+MAX_NEW_TOKENS = int(os.getenv("MAX_NEW_TOKENS", 2048))
+DEVICE = "cuda" if torch.cuda.is_available() else "mps" if torch.backends.mps.is_available() else "cpu"
 
 # --- Model Setup ---
-model_path = "WhiteRabbitNeo/WhiteRabbitNeo-2.5-Qwen-2.5-Coder-7B"
+quant_config = BitsAndBytesConfig(
+    load_in_4bit=LOAD_IN_4BIT,
+    bnb_4bit_compute_dtype=torch.bfloat16,
+    bnb_4bit_use_double_quant=True,
+    bnb_4bit_quant_type="nf4",
+)
 
 model = AutoModelForCausalLM.from_pretrained(
-    model_path,
-    torch_dtype=torch.float16,
-    device_map="auto",
-    load_in_4bit=False,
-    trust_remote_code=False,
+    MODEL_PATH,
+    quantization_config=quant_config if LOAD_IN_4BIT else None,
+    torch_dtype=torch.bfloat16 if DEVICE != "cpu" else torch.float32,
+    device_map="auto" if DEVICE != "cpu" else None,
+    trust_remote_code=True,
 )
 
-tokenizer = AutoTokenizer.from_pretrained(model_path, trust_remote_code=True)
+tokenizer = AutoTokenizer.from_pretrained(MODEL_PATH, trust_remote_code=True)
+
+# --- Generation Function ---
+def generate_code(user_prompt, temperature=0.7, top_p=0.95, max_tokens=1024, top_k=50):
+    if not user_prompt.strip():
+        return "⚠️ Please enter a valid prompt."
 
-def generate_code(user_prompt, temperature=0.75, top_p=1.0, max_tokens=2048, top_k=50):
-    tokens = tokenizer.encode(user_prompt)
-    tokens = torch.LongTensor(tokens).unsqueeze(0).to("cuda")
+    inputs = tokenizer(user_prompt, return_tensors="pt", truncation=True)
+    inputs = {k: v.to(DEVICE) for k, v in inputs.items()}
 
     with torch.no_grad():
         output = model.generate(
-            input_ids=tokens,
-            max_length=len(tokens[0]) + max_tokens,
+            **inputs,
+            max_new_tokens=int(max_tokens),
             do_sample=True,
-            temperature=temperature,
-            top_p=top_p,
-            top_k=top_k,
-            num_return_sequences=1,
+            temperature=float(temperature),
+            top_p=float(top_p),
+            top_k=int(top_k),
             pad_token_id=tokenizer.eos_token_id,
         )
 
-    # Decode only newly generated tokens
-    generated_tokens = output[0][len(tokens[0]):]
-    code_string = tokenizer.decode(generated_tokens, skip_special_tokens=True)
-    # Wrap in python code block for Gradio Markdown rendering
-    return f"```python\n{code_string}\n```"
+    generated_text = tokenizer.decode(output[0], skip_special_tokens=True)
+    new_text = generated_text[len(user_prompt):].strip()
+    safe_code = new_text.replace("```", "`\u200b``")  # Prevent Markdown escape issues
+    return f"```python\n{safe_code}\n```"
 
-# --- Gradio UI ---
+# --- UI ---
 with gr.Blocks(title="Spec Kit Copilot") as demo:
-    with gr.Tab("AI Code Generation"):
-        user_input = gr.Textbox(label="Describe code to generate", lines=4, placeholder="E.g., Python function to sort a list")
-        temperature = gr.Slider(0.0, 1.0, 0.75, label="Temperature")
-        top_p = gr.Slider(0.0, 1.0, 1.0, label="Top-p")
-        max_tokens = gr.Slider(256, 4096, 2048, step=128, label="Max Tokens")
-        top_k = gr.Slider(0, 100, 50, label="Top-k")
-        generate_btn = gr.Button("Generate Code")
-        preview = gr.Markdown()
-        
-        generate_btn.click(
-            fn=generate_code,
-            inputs=[user_input, temperature, top_p, max_tokens, top_k],
-            outputs=preview
-        )
+    gr.Markdown("### 🧠 Spec Kit Copilot — AI Code Generator (Hugging Face Space Edition)")
+    with gr.Row():
+        with gr.Column(scale=2):
+            user_input = gr.Textbox(
+                label="Describe code to generate",
+                lines=4,
+                placeholder="E.g., Python function to parse a JSON file and pretty-print it."
+            )
+            with gr.Row():
+                temperature = gr.Slider(0.0, 1.0, 0.7, label="Temperature")
+                top_p = gr.Slider(0.0, 1.0, 0.95, label="Top-p")
+            with gr.Row():
+                max_tokens = gr.Slider(256, 4096, 1024, step=128, label="Max Tokens")
+                top_k = gr.Slider(0, 100, 50, label="Top-k")
+            generate_btn = gr.Button("🚀 Generate Code")
+        with gr.Column(scale=3):
+            preview = gr.Markdown("")
+
+    generate_btn.click(
+        fn=generate_code,
+        inputs=[user_input, temperature, top_p, max_tokens, top_k],
+        outputs=preview,
+    )
 
 if __name__ == "__main__":
-    demo.launch(server_name="0.0.0.0", server_port=7860)
+    demo.launch(server_name="0.0.0.0", server_port=int(os.getenv("PORT", 7860)))