Update app.py

app.py

@@ -1,8 +1,11 @@
 """
 codenyx.py – CodeNyx AI pair-programmer
-Runs great on HF Free tier (16 GB RAM) or any laptop with 8 GB+ VRAM.
+One-file, zero-config.  If the model is gated, the UI will
+prompt for your HF token once and remember it.
 """
 
+import os
+import sys
 import gradio as gr
 import torch
 from transformers import (
@@ -11,21 +14,46 @@ from transformers import (
     BitsAndBytesConfig,
     TextIteratorStreamer
 )
+from huggingface_hub import login
 from threading import Thread
 
 # ------------------------------------------------------------------
-# 1. Model & branding
+# 1. Branding & constants
 # ------------------------------------------------------------------
-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."
-)
+BOT_NAME      = "CodeNyx"
+MODEL_ID      = "bigcode/starcoder2-3b-instruct"
+SYSTEM_PROMPT = (f"You are {BOT_NAME}, an expert open-source coding assistant. "
+                 "Always provide concise, runnable code snippets with short explanations.")
+
+MAX_NEW_TOK   = 1024
+TEMPERATURE   = 0.2
+TOP_P         = 0.9
+
+# ------------------------------------------------------------------
+# 2. Token helper
+# ------------------------------------------------------------------
+def get_hf_token():
+    """
+    Ask the user interactively once, store in HF_TOKEN env var.
+    If the var already exists, use it silently.
+    """
+    token = os.getenv("HF_TOKEN")
+    if token:
+        return token
+
+    print("🤗  Hugging Face token required for gated/private models.")
+    token = input("Paste your HF token (will be cached in HF_TOKEN): ").strip()
+    if not token:
+        sys.exit("No token provided – aborting.")
+    os.environ["HF_TOKEN"] = token
+    login(token)
+    return token
 
 # ------------------------------------------------------------------
-# 2. 4-bit quantization so CPU-only laptops can run too
+# 3. Load model & tokenizer (with fallback)
 # ------------------------------------------------------------------
+get_hf_token()   # ensures HF_TOKEN is set if needed
+
 bnb_config = BitsAndBytesConfig(
     load_in_4bit=True,
     bnb_4bit_compute_dtype=torch.float16,
@@ -33,21 +61,26 @@ bnb_config = BitsAndBytesConfig(
     bnb_4bit_use_double_quant=True,
 )
 
-tokenizer = AutoTokenizer.from_pretrained(MODEL_ID, trust_remote_code=True)
-tokenizer.pad_token = tokenizer.eos_token  # silence warning
+tokenizer = AutoTokenizer.from_pretrained(
+    MODEL_ID,
+    use_auth_token=os.getenv("HF_TOKEN"),
+    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",
+    use_auth_token=os.getenv("HF_TOKEN"),
     trust_remote_code=True,
 )
 
 # ------------------------------------------------------------------
-# 3. Chat logic
+# 4. Chat logic
 # ------------------------------------------------------------------
 def build_prompt(history, user_input):
-    """Convert Gradio history into model chat template."""
     messages = [{"role": "system", "content": SYSTEM_PROMPT}]
     for human, ai in history:
         messages.append({"role": "user", "content": human})
@@ -58,11 +91,9 @@ def build_prompt(history, user_input):
     )
 
 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)
 
@@ -74,14 +105,13 @@ def bot_turn(history):
 
     gen_kwargs = dict(
         **inputs,
-        max_new_tokens=1024,
-        temperature=0.2,
-        top_p=0.9,
+        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=gen_kwargs).start()
 
     partial = ""
@@ -91,12 +121,12 @@ def bot_turn(history):
         yield history
 
 # ------------------------------------------------------------------
-# 4. Gradio UI
+# 5. Gradio UI
 # ------------------------------------------------------------------
 with gr.Blocks(title=f"{BOT_NAME} – AI Pair-Programmer") as demo:
     gr.Markdown(f"""
     # 🤖 {BOT_NAME} AI Pair-Programmer  
-    *Lightweight 3 B model, 100 % free tier friendly.*  
+    *Runs on CPU or GPU; 3 B parameters, free-tier friendly.*  
     Ask any coding question and get **runnable code + short explanations**.
     """)
 
@@ -104,7 +134,7 @@ with gr.Blocks(title=f"{BOT_NAME} – AI Pair-Programmer") as demo:
 
     with gr.Row():
         msg = gr.Textbox(
-            placeholder="Type your programming question here ...",
+            placeholder="Type your programming question here …",
             lines=2,
             scale=8,
             show_label=False,
@@ -114,7 +144,6 @@ with gr.Blocks(title=f"{BOT_NAME} – AI Pair-Programmer") as demo:
 
     clear_btn = gr.Button("🗑️ Clear")
 
-    # Event wiring
     msg.submit(user_turn, [msg, chatbot], [msg, chatbot], queue=False).then(
         bot_turn, chatbot, chatbot
     )
@@ -124,7 +153,7 @@ with gr.Blocks(title=f"{BOT_NAME} – AI Pair-Programmer") as demo:
     clear_btn.click(lambda: None, None, chatbot, queue=False)
 
 # ------------------------------------------------------------------
-# 5. Launch
+# 6. Launch
 # ------------------------------------------------------------------
 if __name__ == "__main__":
     demo.queue().launch(server_name="0.0.0.0", server_port=7860, share=True)