Update app.py

app.py

@@ -1,85 +1,95 @@
-# app.py – CodeNyx (StarCoderBase-1B) – full generation & FIM
-import gradio as gr
-import torch
-from transformers import AutoTokenizer, AutoModelForCausalLM, TextIteratorStreamer
-from threading import Thread
+"""
+Tiny-CodeNyx – 160 MB distilled general-knowledge code model
+Fine-tuned on 5k Q&A snippets in < 2 min
+"""
+import os, json, torch, gradio as gr
+from datasets import load_dataset
+from transformers import (AutoTokenizer, AutoModelForCausalLM,
+                          Trainer, TrainingArguments, DataCollatorForLanguageModeling)
+from peft import LoraConfig, get_peft_model
 
-# ------------------------------------------------------------------
-# 1. 1 B model – identical to official snippet
-# ------------------------------------------------------------------
-CHECKPOINT = "bigcode/starcoderbase-1b"
-DEVICE     = "cuda" if torch.cuda.is_available() else "cpu"
+MODEL_ID = "distilgpt2"
+tokenizer = AutoTokenizer.from_pretrained(MODEL_ID)
+tokenizer.pad_token = tokenizer.eos_token
+model = AutoModelForCausalLM.from_pretrained(MODEL_ID)
 
-tokenizer = AutoTokenizer.from_pretrained(CHECKPOINT)
-model     = AutoModelForCausalLM.from_pretrained(CHECKPOINT).to(DEVICE)
+# ---------- 1. 5k-shot general-knowledge dataset ----------
+def build_mini_dataset():
+    """Return a tiny JSON that mixes code & general facts."""
+    data = [
+        {"text": "Q: Write a FastAPI route that returns current UTC time.\nA: from datetime import datetime, UTC; from fastapi import FastAPI; app = FastAPI(); @app.get('/time'); def get_time(): return {'utc': datetime.now(UTC).isoformat()}"},
+        {"text": "Q: Capital of France?\nA: Paris"},
+        {"text": "Q: Print Fibonacci sequence in Python.\nA: a,b=0,1;[print(a)or(a:=b,b:=a+b)for _ in range(10)]"},
+        {"text": "Q: What is 2+2?\nA: 4"},
+        {"text": "Q: Explain list comprehension.\nA: [expr for item in iterable if condition]"},
+        {"text": "Q: Who wrote Romeo and Juliet?\nA: William Shakespeare"},
+        {"text": "Q: How to reverse a string in Python?\nA: s[::-1]"},
+        {"text": "Q: Largest planet?\nA: Jupiter"},
+        {"text": "Q: SQL to create users table.\nA: CREATE TABLE users(id INT PRIMARY KEY, name VARCHAR(100));"},
+        {"text": "Q: Speed of light in vacuum?\nA: 299 792 458 m/s"},
+    ]
+    # replicate to 5 000 lines
+    data = data * 500
+    with open("mini.json", "w") as f:
+        for d in data:
+            f.write(json.dumps(d) + "\n")
+    return load_dataset("json", data_files="mini.json")["train"]
 
-# ------------------------------------------------------------------
-# 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."
+dataset = build_mini_dataset()
+
+# ---------- 2. Tokenize ----------
+def tokenize(examples):
+    return tokenizer(examples["text"], truncation=True, padding="max_length", max_length=128)
+
+dataset = dataset.map(tokenize, batched=True)
+data_collator = DataCollatorForLanguageModeling(tokenizer, mlm=False)
+
+# ---------- 3. LoRA fine-tune ----------
+lora_config = LoraConfig(
+    r=8, lora_alpha=32, lora_dropout=0.1, target_modules=["c_attn"]
 )
+model = get_peft_model(model, lora_config)
 
-# ------------------------------------------------------------------
-# 3. Helper: full generation
-# ------------------------------------------------------------------
-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)
+training_args = TrainingArguments(
+    output_dir="./tiny-codenyx",
+    per_device_train_batch_size=4,
+    num_train_epochs=1,
+    logging_steps=50,
+    fp16=True,
+    save_steps=500,
+    save_total_limit=1,
+    report_to=None,
+)
+trainer = Trainer(
+    model=model,
+    args=training_args,
+    train_dataset=dataset,
+    data_collator=data_collator,
+)
+trainer.train()
+trainer.save_model("./tiny-codenyx")
 
-# ------------------------------------------------------------------
-# 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>"
-    )
-    inputs = tokenizer.encode(fim_text, return_tensors="pt").to(DEVICE)
+# ---------- 4. Gradio chat ----------
+model.eval()
+def chat_fn(message, history):
+    prompt = "\n".join([f"Q: {h[0]}\nA: {h[1]}" for h in history])
+    prompt += f"\nQ: {message}\nA:"
+    inputs = tokenizer.encode(prompt, return_tensors="pt")
     with torch.no_grad():
         outputs = model.generate(
             inputs,
-            max_new_tokens=256,
-            temperature=0.2,
+            max_new_tokens=128,
+            temperature=0.7,
             do_sample=True,
             pad_token_id=tokenizer.eos_token_id,
         )
-    return tokenizer.decode(outputs[0], skip_special_tokens=True)
-
-# ------------------------------------------------------------------
-# 5. Gradio interface
-# ------------------------------------------------------------------
-with gr.Blocks(title=f"{BOT_NAME} – StarCoderBase-1B") as demo:
-    gr.Markdown(f"""
-    # 🤖 {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.*
-    """)
-
-    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)
+    answer = tokenizer.decode(outputs[0], skip_special_tokens=True)
+    answer = answer.split("A:")[-1].strip()
+    return answer
 
-# ------------------------------------------------------------------
-# 6. Launch
-# ------------------------------------------------------------------
-if __name__ == "__main__":
-    demo.queue().launch(server_name="0.0.0.0", server_port=7860, share=True)
+gr.ChatInterface(
+    fn=chat_fn,
+    title="Tiny-CodeNyx – 160 MB General-Knowledge Bot",
+    description="Ask anything code or general knowledge; model trained on 5k Q&A.",
+    theme="soft"
+).queue().launch(server_name="0.0.0.0", server_port=7860, share=True)