Update app.py

app.py

@@ -3,17 +3,63 @@ from huggingface_hub import InferenceClient
 import os
 import json
 import random
+from typing import Optional
 
 pipe = None
 
-# Load WPI facts from JSON
-with open("facts.json", "r") as f:
-    WPI_FACTS = json.load(f)
-
-# Fancy CSS (same as before)
-fancy_css = """ ... """  # Keep your previous CSS here
-
-# --- Gompei chatbot response ---
+# ========== Config ==========
+LOCAL_MODEL = os.environ.get("LOCAL_MODEL", "microsoft/Phi-3-mini-4k-instruct")
+
+API_PROVIDER = os.environ.get("API_PROVIDER", "").strip().lower()   # "", "hf", "nebius"
+API_MODEL = os.environ.get("API_MODEL", "HuggingFaceH4/zephyr-7b-beta")
+
+NEBIUS_API_KEY = os.environ.get("NEBIUS_API_KEY")
+NEBIUS_MODEL = os.environ.get("NEBIUS_MODEL", "gpt-oss-20b")
+NEBIUS_BASE_URL = os.environ.get("NEBIUS_BASE_URL", "https://api.studio.nebius.ai/v1")
+# ===========================
+
+# Facts + CSS fallbacks
+FACTS_PATH = "facts.json"
+DEFAULT_FACTS = [{"text": "WPI was founded in 1865 by John Boynton and Ichabod Washburn."}]
+try:
+    with open(FACTS_PATH, "r") as f:
+        WPI_FACTS = json.load(f)
+    if not isinstance(WPI_FACTS, list) or not WPI_FACTS:
+        WPI_FACTS = DEFAULT_FACTS
+except Exception:
+    WPI_FACTS = DEFAULT_FACTS
+
+fancy_css = """/* fallback if your CSS file isn't ready */ #title { text-align:center; }"""
+
+def _extract_hf_token(hf_token_obj: Optional[object]) -> Optional[str]:
+    if hf_token_obj:
+        if isinstance(hf_token_obj, str) and hf_token_obj.strip():
+            return hf_token_obj.strip()
+        for attr in ("token", "access_token"):
+            try:
+                val = getattr(hf_token_obj, attr, None)
+                if isinstance(val, str) and val.strip():
+                    return val.strip()
+            except Exception:
+                pass
+        try:
+            if hasattr(hf_token_obj, "get"):
+                val = hf_token_obj.get("token") or hf_token_obj.get("access_token")
+                if isinstance(val, str) and val.strip():
+                    return val.strip()
+        except Exception:
+            pass
+    env_val = os.environ.get("HF_TOKEN")
+    if isinstance(env_val, str) and env_val.strip():
+        return env_val.strip()
+    return None
+
+def _resolve_provider():
+    if API_PROVIDER in ("hf", "nebius"):
+        return API_PROVIDER
+    return "nebius" if NEBIUS_API_KEY else "hf"
+
+# ---- Core chat handler (unchanged logic) ----
 def respond(
     message,
     history: list[dict[str, str]],
@@ -22,16 +68,13 @@ def respond(
     temperature,
     top_p,
     use_local_model: bool,
-    hf_token: gr.OAuthToken = None,
+    hf_token: Optional[object] = None,
 ):
     global pipe
 
-    # Pick a random WPI fact
     fact = random.choice(WPI_FACTS)["text"]
-
     messages = [{"role": "system", "content": system_message}]
     messages.extend(history)
-    # Append the user message + random fact for context
     messages.append({"role": "user", "content": f"{message}\n\nFun fact: {fact}"})
 
     response = ""
@@ -40,10 +83,8 @@ def respond(
         print("[MODE] local")
         from transformers import pipeline
         if pipe is None:
-            pipe = pipeline("text-generation", model="microsoft/Phi-3-mini-4k-instruct")
-
+            pipe = pipeline("text-generation", model=LOCAL_MODEL)
         prompt = "\n".join([f"{m['role']}: {m['content']}" for m in messages])
-
         outputs = pipe(
             prompt,
             max_new_tokens=max_tokens,
@@ -51,78 +92,115 @@ def respond(
             temperature=temperature,
             top_p=top_p,
         )
-
         response = outputs[0]["generated_text"][len(prompt):]
         yield response.strip()
+        return
 
-    else:
-        print("[MODE] api")
-        token_value = None
-        if hf_token and getattr(hf_token, "token", None):
-            token_value = hf_token.token
-        elif os.environ.get("HF_TOKEN"):
-            token_value = os.environ.get("HF_TOKEN")
+    provider = _resolve_provider()
 
-        if not token_value:
-            yield "⚠️ Please log in with your Hugging Face account or set HF_TOKEN in environment."
+    if provider == "nebius":
+        print(f"[MODE] api | provider=nebius model={NEBIUS_MODEL}")
+        if not NEBIUS_API_KEY:
+            yield ("⚠️ Missing NEBIUS_API_KEY. Set it or switch to HF by setting API_PROVIDER=hf and providing HF_TOKEN.")
             return
-
-        client = InferenceClient(token=token_value, model="openai/gpt-oss-20b")
-
-        for chunk in client.chat_completion(
-            messages,
-            max_tokens=max_tokens,
-            stream=True,
+        client = InferenceClient(token=NEBIUS_API_KEY, base_url=NEBIUS_BASE_URL)
+        try:
+            for chunk in client.chat_completion(
+                messages=messages,
+                max_tokens=max_tokens,
+                stream=True,
+                temperature=temperature,
+                top_p=top_p,
+                model=NEBIUS_MODEL,
+            ):
+                choices = getattr(chunk, "choices", [])
+                token_text = ""
+                if choices and getattr(choices[0].delta, "content", None):
+                    token_text = choices[0].delta.content
+                response += token_text
+                yield response
+        except Exception as e:
+            if "401" in str(e) or "Unauthorized" in str(e):
+                yield "⚠️ Nebius auth failed. Check NEBIUS_API_KEY and NEBIUS_MODEL."
+            else:
+                yield f"⚠️ Nebius API error: {e}"
+        return
+
+    # HF provider via text_generation (no strict chat perms)
+    print(f"[MODE] api | provider=hf model={API_MODEL}")
+    token_value = _extract_hf_token(hf_token)
+    if not token_value:
+        yield "⚠️ Please log in (Login button) or set HF_TOKEN in environment."
+        return
+    client = InferenceClient(model=API_MODEL, token=token_value)
+
+    prompt = "\n".join([f"{m['role']}: {m['content']}" for m in messages])
+    try:
+        stream = client.text_generation(
+            prompt,
+            max_new_tokens=max_tokens,
             temperature=temperature,
             top_p=top_p,
-        ):
-            choices = chunk.choices
-            token = ""
-            if len(choices) and choices[0].delta.content:
-                token = choices[0].delta.content
-            response += token
+            stream=True,
+            details=False,
+            return_full_text=False,
+        )
+        for out in stream:
+            try:
+                token_text = getattr(out, "token", None)
+                token_text = token_text.text if token_text else (out if isinstance(out, str) else "")
+            except Exception:
+                token_text = str(out) if out else ""
+            response += token_text
             yield response
+    except Exception as e:
+        if "401" in str(e) or "Unauthorized" in str(e):
+            yield "⚠️ Hugging Face auth failed. Ensure HF_TOKEN or log in via the button."
+        else:
+            yield f"⚠️ HF Inference error: {e}"
+
+# ---- Build UI only when asked ----
+def create_demo(enable_oauth: bool = True):
+    with gr.Blocks(css=fancy_css) as demo:
+        with gr.Row():
+            gr.Markdown("<h1 id='title'>🐐 Chat with Gompei</h1>")
+            token_input = gr.LoginButton() if enable_oauth else gr.State(value=None)
+
+        gr.ChatInterface(
+            fn=respond,
+            additional_inputs=[
+                gr.Textbox(
+                    value="You are Gompei the Goat, WPI's mascot. Answer questions with fun goat-like personality and real WPI facts.",
+                    label="System message",
+                ),
+                gr.Slider(minimum=1, maximum=1024, value=256, step=1, label="Max new tokens"),
+                gr.Slider(minimum=0.1, maximum=2.0, value=0.7, step=0.1, label="Temperature"),
+                gr.Slider(minimum=0.1, maximum=1.0, value=0.95, step=0.05, label="Top-p (nucleus sampling)"),
+                gr.Checkbox(label="Use Local Model", value=False),
+                token_input,  # LoginButton or a dummy State(None) to keep signature aligned
+            ],
+            type="messages",
+            examples=[
+                [
+                    "Where is WPI located?",
+                    "You are Gompei the Goat, WPI's mascot. Answer questions with fun goat-like personality and real WPI facts.",
+                    128, 0.7, 0.95, False, None
+                ],
+                [
+                    "Who founded WPI?",
+                    "You are Gompei the Goat, WPI's mascot. Answer questions with fun goat-like personality and real WPI facts.",
+                    128, 0.7, 0.95, False, None
+                ],
+            ],
+        )
+    return demo
 
-# --- Chat Interface ---
-chatbot = gr.ChatInterface(
-    fn=respond,
-    additional_inputs=[
-        gr.Textbox(
-            value="You are Gompei the Goat, WPI's mascot. Answer questions with fun goat-like personality and real WPI facts.",
-            label="System message",
-        ),
-        gr.Slider(minimum=1, maximum=1024, value=256, step=1, label="Max new tokens"),
-        gr.Slider(minimum=0.1, maximum=2.0, value=0.7, step=0.1, label="Temperature"),
-        gr.Slider(minimum=0.1, maximum=1.0, value=0.95, step=0.05, label="Top-p (nucleus sampling)"),
-        gr.Checkbox(label="Use Local Model", value=False),
-    ],
-    type="messages",
-    examples=[
-        [
-            "Where is WPI located?",
-            "You are Gompei the Goat, WPI's mascot. Answer questions with fun goat-like personality and real WPI facts.",
-            128,
-            0.7,
-            0.95,
-            False
-        ],
-        [
-            "Who founded WPI?",
-            "You are Gompei the Goat, WPI's mascot. Answer questions with fun goat-like personality and real WPI facts.",
-            128,
-            0.7,
-            0.95,
-            False
-        ],
-    ],
-)
-
-# --- Blocks layout ---
-with gr.Blocks(css=fancy_css) as demo:
-    with gr.Row():
-        gr.Markdown("<h1 id='title'>🐐 Chat with Gompei</h1>")
-        gr.LoginButton()
-    chatbot.render()
+# Create demo automatically unless we're in CI/tests
+if os.environ.get("SKIP_UI_ON_IMPORT") != "1":
+    demo = create_demo(enable_oauth=True)
 
 if __name__ == "__main__":
+    # If not created above (e.g., when SKIP_UI_ON_IMPORT=1 locally), create now
+    if "demo" not in globals():
+        demo = create_demo(enable_oauth=True)
     demo.launch(server_name="0.0.0.0", server_port=7860)