Update app.py

app.py

@@ -1,127 +1,85 @@
-"""
-CodeNyx – 100 % ready-to-run Space
-- Model loads at startup, no button needed
-- 3 B StarCoder2 4-bit → < 8 GB RAM
-- Answers generated directly from The Stack v2 (code-only corpus)
-"""
-
-import os
+# app.py – CodeNyx (StarCoderBase-1B) – full generation & FIM
 import gradio as gr
 import torch
-from transformers import (
-    AutoTokenizer,
-    AutoModelForCausalLM,
-    BitsAndBytesConfig,
-    TextIteratorStreamer
-)
+from transformers import AutoTokenizer, AutoModelForCausalLM, TextIteratorStreamer
 from threading import Thread
 
-BOT_NAME = "CodeNyx"
-MODEL_ID = "bigcode/starcoder2-3b"
-
-MAX_NEW_TOK = 1024
-TEMPERATURE = 0.2
-TOP_P = 0.9
-
-SYSTEM_PROMPT = (
-    f"You are {BOT_NAME}, an expert open-source coding assistant trained on "
-    "The Stack v2. Always return concise, runnable code snippets with brief explanations."
-)
-
 # ------------------------------------------------------------------
-# 1. Startup model loader (runs once)
+# 1. 1 B model – identical to official snippet
 # ------------------------------------------------------------------
-print("🚀 CodeNyx – loading model …")
-bnb_config = BitsAndBytesConfig(
-    load_in_4bit=True,
-    bnb_4bit_compute_dtype=torch.float16,
-    bnb_4bit_quant_type="nf4",
-    bnb_4bit_use_double_quant=True,
-)
+CHECKPOINT = "bigcode/starcoderbase-1b"
+DEVICE     = "cuda" if torch.cuda.is_available() else "cpu"
 
-tokenizer = AutoTokenizer.from_pretrained(
-    MODEL_ID,
-    trust_remote_code=True,
-)
-if tokenizer.pad_token is None:
-    tokenizer.pad_token = tokenizer.eos_token
+tokenizer = AutoTokenizer.from_pretrained(CHECKPOINT)
+model     = AutoModelForCausalLM.from_pretrained(CHECKPOINT).to(DEVICE)
 
-model = AutoModelForCausalLM.from_pretrained(
-    MODEL_ID,
-    quantization_config=bnb_config,
-    device_map="auto",
-    trust_remote_code=True,
+# ------------------------------------------------------------------
+# 2. Branding
+# ------------------------------------------------------------------
+BOT_NAME = "CodeNyx"
+SYSTEM   = (
+    f"You are {BOT_NAME}, an expert coding assistant trained on The Stack v1.2. "
+    "Return only complete, runnable code with a short comment."
 )
-print("✅ CodeNyx model and tokenizer loaded.")
 
 # ------------------------------------------------------------------
-# 2. Chat helpers
+# 3. Helper: full generation
 # ------------------------------------------------------------------
-def build_prompt(history, user_input):
-    """Turn chat history into the model’s prompt format."""
-    messages = [{"role": "system", "content": SYSTEM_PROMPT}]
-    for human, ai in history:
-        messages += [{"role": "user", "content": human},
-                     {"role": "assistant", "content": ai}]
-    messages.append({"role": "user", "content": user_input})
-    return tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)
-
-def user_turn(user_message, history):
-    return "", history + [[user_message, ""]]
-
-def bot_turn(history):
-    prompt = build_prompt(history[:-1], history[-1][0])
-    inputs = tokenizer(prompt, return_tensors="pt").to(model.device)
+def full_generation(prompt: str):
+    inputs = tokenizer.encode(prompt, return_tensors="pt").to(DEVICE)
+    with torch.no_grad():
+        outputs = model.generate(
+            inputs,
+            max_new_tokens=512,
+            temperature=0.2,
+            do_sample=True,
+            pad_token_id=tokenizer.eos_token_id,
+        )
+    return tokenizer.decode(outputs[0], skip_special_tokens=True)
 
-    streamer = TextIteratorStreamer(tokenizer, skip_prompt=True, skip_special_tokens=True)
-    gen_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,
+# ------------------------------------------------------------------
+# 4. Helper: fill-in-the-middle (FIM)
+# ------------------------------------------------------------------
+def fim_generation(prefix: str, suffix: str):
+    fim_text = (
+        f"<fim_prefix>{prefix}<fim_suffix>{suffix}<fim_middle>"
     )
-    Thread(target=model.generate, kwargs=gen_kwargs).start()
-
-    partial = ""
-    for new_text in streamer:
-        partial += new_text
-        history[-1][1] = partial
-        yield history
+    inputs = tokenizer.encode(fim_text, return_tensors="pt").to(DEVICE)
+    with torch.no_grad():
+        outputs = model.generate(
+            inputs,
+            max_new_tokens=256,
+            temperature=0.2,
+            do_sample=True,
+            pad_token_id=tokenizer.eos_token_id,
+        )
+    return tokenizer.decode(outputs[0], skip_special_tokens=True)
 
 # ------------------------------------------------------------------
-# 3. Gradio UI
+# 5. Gradio interface
 # ------------------------------------------------------------------
-with gr.Blocks(title=f"{BOT_NAME} – AI Pair-Programmer") as demo:
+with gr.Blocks(title=f"{BOT_NAME} – StarCoderBase-1B") as demo:
     gr.Markdown(f"""
-    # 🤖 {BOT_NAME} – AI Pair-Programmer
-    Trained on **The Stack v2** (4 T tokens, permissive licences).  
-    Ask any coding question and receive **runnable code + short explanations**.
+    # 🤖 {BOT_NAME} – powered by StarCoderBase-1B (The Stack v1.2)
+    *Ask for full code or let the model **fill-in-the-middle** of any snippet.*
     """)
-    chatbot = gr.Chatbot(height=500, label=f"{BOT_NAME} Chat")
-    with gr.Row():
-        msg = gr.Textbox(
-            placeholder="Type your programming question here …",
-            lines=2,
-            scale=8,
-            show_label=False,
-            container=False,
-        )
-        send_btn = gr.Button("Send", scale=1, variant="primary")
-    clear_btn = gr.Button("🗑️ Clear")
 
-    msg.submit(user_turn, [msg, chatbot], [msg, chatbot], queue=False).then(
-        bot_turn, chatbot, chatbot
-    )
-    send_btn.click(user_turn, [msg, chatbot], [msg, chatbot], queue=False).then(
-        bot_turn, chatbot, chatbot
-    )
-    clear_btn.click(lambda: None, None, chatbot, queue=False)
+    with gr.Tab("Full Generation"):
+        prompt_in   = gr.Textbox(label="Prompt", lines=3, placeholder="def fibonacci(n):")
+        full_out    = gr.Code(label="Generated Code", language="python")
+        gen_btn     = gr.Button("Generate")
+        gen_btn.click(full_generation, prompt_in, full_out)
+
+    with gr.Tab("Fill-in-the-Middle"):
+        with gr.Row():
+            prefix_in = gr.Textbox(label="Prefix", lines=3, placeholder="def fibonacci(n):\n    ")
+            suffix_in = gr.Textbox(label="Suffix", lines=3, placeholder="\n    return result")
+        fim_out = gr.Code(label="Completed Code", language="python")
+        fim_btn = gr.Button("Complete")
+        fim_btn.click(fim_generation, [prefix_in, suffix_in], fim_out)
 
 # ------------------------------------------------------------------
-# 4. Launch
+# 6. Launch
 # ------------------------------------------------------------------
 if __name__ == "__main__":
     demo.queue().launch(server_name="0.0.0.0", server_port=7860, share=True)