Update app.py

app.py

@@ -1,11 +1,9 @@
 """
-codenyx.py – CodeNyx AI pair-programmer
-One-file, zero-config.  If the model is gated, the UI will
-prompt for your HF token once and remember it.
+CodeNyx – zero-config Gradio chatbot
+Auto-handles Hugging Face tokens without stdin prompts.
 """
 
 import os
-import sys
 import gradio as gr
 import torch
 from transformers import (
@@ -17,11 +15,8 @@ from transformers import (
 from huggingface_hub import login
 from threading import Thread
 
-# ------------------------------------------------------------------
-# 1. Branding & constants
-# ------------------------------------------------------------------
-BOT_NAME      = "CodeNyx"
 MODEL_ID      = "bigcode/starcoder2-3b-instruct"
+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.")
 
@@ -30,55 +25,57 @@ TEMPERATURE   = 0.2
 TOP_P         = 0.9
 
 # ------------------------------------------------------------------
-# 2. Token helper
+# 1. Token helper (Gradio popup instead of stdin)
 # ------------------------------------------------------------------
-def get_hf_token():
+def ensure_token(token_value):
     """
-    Ask the user interactively once, store in HF_TOKEN env var.
-    If the var already exists, use it silently.
+    token_value comes from the Gradio UI the first time.
+    We cache it in environment variable HF_TOKEN and login once.
+    Returns True on success.
     """
-    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()
+    token = token_value.strip()
     if not token:
-        sys.exit("No token provided – aborting.")
+        return False
     os.environ["HF_TOKEN"] = token
     login(token)
-    return token
+    return True
 
 # ------------------------------------------------------------------
-# 3. Load model & tokenizer (with fallback)
+# 2. Lazy model loader (once token is ready)
 # ------------------------------------------------------------------
-get_hf_token()   # ensures HF_TOKEN is set if needed
-
-bnb_config = BitsAndBytesConfig(
-    load_in_4bit=True,
-    bnb_4bit_compute_dtype=torch.float16,
-    bnb_4bit_quant_type="nf4",
-    bnb_4bit_use_double_quant=True,
-)
+model, tokenizer = None, None
+
+def load_model():
+    global model, tokenizer
+    if model is not None:
+        return True  # already loaded
+
+    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,
-    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,
-)
+    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,
+    )
+    return True
 
 # ------------------------------------------------------------------
-# 4. Chat logic
+# 3. Chat logic
 # ------------------------------------------------------------------
 def build_prompt(history, user_input):
     messages = [{"role": "system", "content": SYSTEM_PROMPT}]
@@ -121,39 +118,57 @@ def bot_turn(history):
         yield history
 
 # ------------------------------------------------------------------
-# 5. Gradio UI
+# 4. Gradio UI flow
 # ------------------------------------------------------------------
 with gr.Blocks(title=f"{BOT_NAME} – AI Pair-Programmer") as demo:
     gr.Markdown(f"""
-    # 🤖 {BOT_NAME} AI Pair-Programmer  
-    *Runs on CPU or GPU; 3 B parameters, free-tier friendly.*  
-    Ask any coding question and get **runnable code + short explanations**.
+    # 🤖 {BOT_NAME} – AI Pair-Programmer  
+    *3 B params, 100 % free-tier friendly.*  
+    Paste your Hugging Face token **once** if asked, then chat away.
     """)
 
-    chatbot = gr.Chatbot(height=500)
-
-    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")
-
-    clear_btn = gr.Button("🗑️ Clear")
-
-    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_btn.click(lambda: None, None, chatbot, queue=False)
-
-# ------------------------------------------------------------------
-# 6. Launch
-# ------------------------------------------------------------------
-if __name__ == "__main__":
-    demo.queue().launch(server_name="0.0.0.0", server_port=7860, share=True)
+    token_box  = gr.Textbox(label="🤗 Hugging Face Token (only first time)", type="password")
+    load_btn   = gr.Button("Authorize")
+    status_lbl = gr.Label(value="Waiting for token …")
+
+    # --- main chat controls (hidden until token OK) ---
+    with gr.Column(visible=False) as chat_col:
+        chatbot = gr.Chatbot(height=450)
+        with gr.Row():
+            msg = gr.Textbox(
+                placeholder="Ask me to write / debug / explain code …",
+                lines=2,
+                scale=8,
+                show_label=False,
+                container=False,
+            )
+            send_btn = gr.Button("Send", scale=1, variant="primary")
+        clear_btn = gr.Button("🗑️ Clear")
+
+    # ------------------------------------------------------------------
+    # 5. Event wiring
+    # ------------------------------------------------------------------
+    def _auth(token):
+        ok = ensure_token(token)
+        if ok:
+            try:
+                load_model()
+                return gr.update(visible=False), gr.update(visible=True), "✅ Ready! Start coding."
+            except Exception as e:
+                return gr.update(visible=True), gr.update(visible=False), f"❌ Error: {e}"
+        else:
+            return gr.update(visible=True), gr.update(visible=False), "❌ Invalid token."
+
+    load_btn.click(_auth, token_box, [token_box, chat_col, status_lbl])
+
+    def _send(user_msg, hist):
+        return user_turn(user_msg, hist)
+
+    def _bot(hist):
+        yield from bot_turn(hist)
+
+    msg.submit(_send, [msg, chatbot], [msg, chatbot]).then(_bot, chatbot, chatbot)
+    send_btn.click(_send, [msg, chatbot], [msg, chatbot]).then(_bot, chatbot, chatbot)
+    clear_btn.click(lambda: None, None, chatbot)
+
+demo.queue().launch(server_name="0.0.0.0", server_port=7860, share=True)