Create app.py

app.py

@@ -0,0 +1,125 @@
+"""
+codenyx.py – CodeNyx Coding Assistant
+Powered by StarCoder2-15B-Instruct (fine-tuned on The Stack v2)
+"""
+import gradio as gr
+import torch
+from transformers import (
+    AutoTokenizer,
+    AutoModelForCausalLM,
+    BitsAndBytesConfig,
+    TextIteratorStreamer
+)
+from threading import Thread
+
+# ------------------------------------------------------------------
+# 1. Branding & defaults
+# ------------------------------------------------------------------
+BOT_NAME      = "CodeNyx"
+SYSTEM_PROMPT = (f"You are {BOT_NAME}, an expert open-source coding assistant. "
+                 "Always produce concise, runnable code snippets with explanations.")
+
+MODEL_ID      = "bigcode/starcoder2-15b-instruct-v0.1"
+MAX_NEW_TOK   = 2048
+TEMPERATURE   = 0.2
+TOP_P         = 0.9
+
+# ------------------------------------------------------------------
+# 2. Load model & tokenizer (4-bit quantized by default)
+# ------------------------------------------------------------------
+bnb_config = BitsAndBytesConfig(
+    load_in_4bit=True,
+    bnb_4bit_compute_dtype=torch.float16,
+    bnb_4bit_quant_type="nf4",
+    bnb_4bit_use_double_quant=True
+)
+
+tokenizer = AutoTokenizer.from_pretrained(MODEL_ID, trust_remote_code=True)
+if tokenizer.pad_token is None:
+    tokenizer.pad_token = tokenizer.eos_token
+
+model = AutoModelForCausalLM.from_pretrained(
+    MODEL_ID,
+    quantization_config=bnb_config,
+    device_map="auto",
+    trust_remote_code=True
+)
+
+# ------------------------------------------------------------------
+# 3. Chat helper
+# ------------------------------------------------------------------
+def build_prompt(history, user_input):
+    """Turn Gradio history into the model’s chat format."""
+    messages = [{"role": "system", "content": SYSTEM_PROMPT}]
+    for h in history:
+        messages.append({"role": "user", "content": h[0]})
+        messages.append({"role": "assistant", "content": h[1]})
+    messages.append({"role": "user", "content": user_input})
+    return tokenizer.apply_chat_template(
+        messages, tokenize=False, add_generation_prompt=True
+    )
+
+def chat_fn(message, history):
+    prompt = build_prompt(history, message)
+    inputs = tokenizer(prompt, return_tensors="pt").to(model.device)
+
+    streamer = TextIteratorStreamer(
+        tokenizer,
+        skip_prompt=True,
+        skip_special_tokens=True
+    )
+
+    generation_kwargs = dict(
+        **inputs,
+        max_new_tokens=MAX_NEW_TOK,
+        temperature=TEMPERATURE,
+        top_p=TOP_P,
+        do_sample=True,
+        pad_token_id=tokenizer.eos_token_id,
+        streamer=streamer
+    )
+    Thread(target=model.generate, kwargs=generation_kwargs).start()
+
+    partial = ""
+    for new_text in streamer:
+        partial += new_text
+        yield partial
+
+# ------------------------------------------------------------------
+# 4. Gradio UI
+# ------------------------------------------------------------------
+with gr.Blocks(title=f"{BOT_NAME} – AI Pair-Programmer") as demo:
+    gr.Markdown(f"""
+    # 🤖 {BOT_NAME}
+    *An open-source coding assistant trained on The Stack v2 – 15 B parameters, permissive licenses only.*  
+    Type any programming question, and {BOT_NAME} will return **runnable code + concise explanations**.
+    """)
+
+    chatbot = gr.Chatbot(
+        label=f"{BOT_NAME} Chat",
+        height=600,
+        avatar_images=(None, "https://huggingface.co/spaces/bigcode/starcoder-logo/resolve/main/logo.png")
+    )
+
+    msg = gr.Textbox(
+        placeholder="Ask me to write, debug, or explain code...",
+        lines=2,
+        show_label=False,
+        container=False
+    )
+
+    clear = gr.Button("🗑️ Clear history")
+
+    def user(user_message, history):
+        return "", history + [[user_message, ""]]
+
+    msg.submit(user, [msg, chatbot], [msg, chatbot], queue=False).then(
+        chat_fn, [msg, chatbot], chatbot
+    )
+    clear.click(lambda: None, None, chatbot, queue=False)
+
+# ------------------------------------------------------------------
+# 5. Launch
+# ------------------------------------------------------------------
+if __name__ == "__main__":
+    demo.queue().launch(server_name="0.0.0.0", server_port=7860, share=True)