Update app.py

app.py

@@ -1,16 +1,59 @@
 
 import gradio as gr
+import torch
+from transformers import AutoModelForCausalLM, AutoTokenizer
 
-def predict(text: str) -> str:
-    return text[::-1]  # toy example: reverse input
+# --- Model Setup ---
+model_path = "WhiteRabbitNeo/WhiteRabbitNeo-2.5-Qwen-2.5-Coder-7B"
 
-with gr.Blocks() as demo:
-    gr.Markdown("## Pure Gradio App")
-    with gr.Row():
-        inp = gr.Textbox(label="Input")
-        out = gr.Textbox(label="Output")
-    btn = gr.Button("Run")
-    btn.click(fn=predict, inputs=inp, outputs=out)
+model = AutoModelForCausalLM.from_pretrained(
+    model_path,
+    torch_dtype=torch.float16,
+    device_map="auto",
+    load_in_4bit=False,
+    trust_remote_code=False,
+)
+
+tokenizer = AutoTokenizer.from_pretrained(model_path, trust_remote_code=True)
+
+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")
+
+    with torch.no_grad():
+        output = model.generate(
+            input_ids=tokens,
+            max_length=len(tokens[0]) + max_tokens,
+            do_sample=True,
+            temperature=temperature,
+            top_p=top_p,
+            top_k=top_k,
+            num_return_sequences=1,
+            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```"
+
+# --- Gradio 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
+        )
 
 if __name__ == "__main__":
-    demo.launch()
+    demo.launch(server_name="0.0.0.0", server_port=7860)