Spaces:
Runtime error
Runtime error
File size: 3,102 Bytes
b2ce241 8d6b3ac b2ce241 31d8051 8d6b3ac b2ce241 31d8051 8d6b3ac b2ce241 8d6b3ac b2ce241 8d6b3ac b2ce241 8d6b3ac b2ce241 8d6b3ac b2ce241 8d6b3ac b2ce241 8d6b3ac b2ce241 8d6b3ac b2ce241 31d8051 b2ce241 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 |
import os
import torch
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 ---
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,
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)
# --- 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."
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(
**inputs,
max_new_tokens=int(max_tokens),
do_sample=True,
temperature=float(temperature),
top_p=float(top_p),
top_k=int(top_k),
pad_token_id=tokenizer.eos_token_id,
)
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```"
# --- UI ---
with gr.Blocks(title="Spec Kit Copilot") as demo:
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=int(os.getenv("PORT", 7860)))
|