Update app.py

app.py

@@ -1,7 +1,8 @@
 """
-codenyx.py – CodeNyx Coding Assistant
-Powered by StarCoder2-15B-Instruct (fine-tuned on The Stack v2)
+codenyx.py – CodeNyx AI pair-programmer
+Runs great on HF Free tier (16 GB RAM) or any laptop with 8 GB+ VRAM.
 """
+
 import gradio as gr
 import torch
 from transformers import (
@@ -13,110 +14,114 @@ from transformers import (
 from threading import Thread
 
 # ------------------------------------------------------------------
-# 1. Branding & defaults
+# 1. Model & branding
 # ------------------------------------------------------------------
-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
+MODEL_ID = "bigcode/starcoder2-3b-instruct-v0.1"  # 3 B params, fits 16 GB
+BOT_NAME = "CodeNyx"
+SYSTEM_PROMPT = (
+    f"You are {BOT_NAME}, an expert open-source coding assistant. "
+    "Always provide concise, runnable code snippets with short explanations."
+)
 
 # ------------------------------------------------------------------
-# 2. Load model & tokenizer (4-bit quantized by default)
+# 2. 4-bit quantization so CPU-only laptops can run too
 # ------------------------------------------------------------------
 bnb_config = BitsAndBytesConfig(
     load_in_4bit=True,
     bnb_4bit_compute_dtype=torch.float16,
     bnb_4bit_quant_type="nf4",
-    bnb_4bit_use_double_quant=True
+    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
+tokenizer.pad_token = tokenizer.eos_token  # silence warning
 
 model = AutoModelForCausalLM.from_pretrained(
     MODEL_ID,
     quantization_config=bnb_config,
     device_map="auto",
-    trust_remote_code=True
+    trust_remote_code=True,
 )
 
 # ------------------------------------------------------------------
-# 3. Chat helper
+# 3. Chat logic
 # ------------------------------------------------------------------
 def build_prompt(history, user_input):
-    """Turn Gradio history into the model’s chat format."""
+    """Convert Gradio history into model chat template."""
     messages = [{"role": "system", "content": SYSTEM_PROMPT}]
-    for h in history:
-        messages.append({"role": "user", "content": h[0]})
-        messages.append({"role": "assistant", "content": h[1]})
+    for human, ai in history:
+        messages.append({"role": "user", "content": human})
+        messages.append({"role": "assistant", "content": ai})
     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)
+def user_turn(user_message, history):
+    """Add user message to history."""
+    return "", history + [[user_message, ""]]
+
+def bot_turn(history):
+    """Generate assistant reply in streaming mode."""
+    prompt = build_prompt(history[:-1], history[-1][0])
     inputs = tokenizer(prompt, return_tensors="pt").to(model.device)
 
     streamer = TextIteratorStreamer(
         tokenizer,
         skip_prompt=True,
-        skip_special_tokens=True
+        skip_special_tokens=True,
     )
 
-    generation_kwargs = dict(
+    gen_kwargs = dict(
         **inputs,
-        max_new_tokens=MAX_NEW_TOK,
-        temperature=TEMPERATURE,
-        top_p=TOP_P,
+        max_new_tokens=1024,
+        temperature=0.2,
+        top_p=0.9,
         do_sample=True,
         pad_token_id=tokenizer.eos_token_id,
-        streamer=streamer
+        streamer=streamer,
     )
-    Thread(target=model.generate, kwargs=generation_kwargs).start()
+
+    Thread(target=model.generate, kwargs=gen_kwargs).start()
 
     partial = ""
     for new_text in streamer:
         partial += new_text
-        yield partial
+        history[-1][1] = partial
+        yield history
 
 # ------------------------------------------------------------------
 # 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**.
+    # 🤖 {BOT_NAME} AI Pair-Programmer  
+    *Lightweight 3 B model, 100 % free tier friendly.*  
+    Ask any coding question and get **runnable code + short 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
-    )
+    chatbot = gr.Chatbot(height=500)
 
-    clear = gr.Button("🗑️ Clear history")
+    with gr.Row():
+        msg = gr.Textbox(
+            placeholder="Type your programming question here ...",
+            lines=2,
+            scale=8,
+            show_label=False,
+            container=False,
+        )
+        submit_btn = gr.Button("Send", scale=1, variant="primary")
 
-    def user(user_message, history):
-        return "", history + [[user_message, ""]]
+    clear_btn = gr.Button("🗑️ Clear")
 
-    msg.submit(user, [msg, chatbot], [msg, chatbot], queue=False).then(
-        chat_fn, [msg, chatbot], chatbot
+    # Event wiring
+    msg.submit(user_turn, [msg, chatbot], [msg, chatbot], queue=False).then(
+        bot_turn, chatbot, chatbot
+    )
+    submit_btn.click(user_turn, [msg, chatbot], [msg, chatbot], queue=False).then(
+        bot_turn, chatbot, chatbot
     )
-    clear.click(lambda: None, None, chatbot, queue=False)
+    clear_btn.click(lambda: None, None, chatbot, queue=False)
 
 # ------------------------------------------------------------------
 # 5. Launch