Update app.py

app.py

@@ -1,98 +1,346 @@
-import spaces
-import gradio as gr
+import os
+import gc
 import torch
-from transformers import AutoModelForCausalLM, AutoTokenizer
+import gradio as gr
+
+from transformers import AutoTokenizer, AutoModelForCausalLM
 
 MODEL_ID = "devoppro/FastLLM"
-TOKENIZER_ID = "Qwen/Qwen2.5-0.5B"  # FastLLM reuses the Qwen2.5 BPE vocab
 
-tokenizer = AutoTokenizer.from_pretrained(TOKENIZER_ID)
-model = AutoModelForCausalLM.from_pretrained(
+# ---------------------------------------------------------
+# Environment
+# ---------------------------------------------------------
+
+os.environ.setdefault("HF_HOME", "/tmp/huggingface")
+os.environ.setdefault("TOKENIZERS_PARALLELISM", "false")
+
+DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
+
+print("=" * 60)
+print("FastLLM")
+print("=" * 60)
+print(f"Model:  {MODEL_ID}")
+print(f"Device: {DEVICE}")
+
+# ---------------------------------------------------------
+# Load tokenizer
+# ---------------------------------------------------------
+
+print("Loading tokenizer...")
+
+tokenizer = AutoTokenizer.from_pretrained(
     MODEL_ID,
-    torch_dtype=torch.float16,
     trust_remote_code=True,
 )
-model.to("cuda")
+
+if tokenizer.pad_token is None:
+    tokenizer.pad_token = tokenizer.eos_token
+
+print("Tokenizer loaded.")
+
+# ---------------------------------------------------------
+# Load model
+# ---------------------------------------------------------
+
+print("Loading FastLLM...")
+
+model_kwargs = {
+    "trust_remote_code": True,
+    "low_cpu_mem_usage": True,
+}
+
+if DEVICE == "cuda":
+    model_kwargs["torch_dtype"] = torch.float16
+else:
+    model_kwargs["torch_dtype"] = torch.float32
+
+model = AutoModelForCausalLM.from_pretrained(
+    MODEL_ID,
+    **model_kwargs,
+)
+
 model.eval()
+model.to(DEVICE)
+
+print("FastLLM loaded successfully.")
+print("=" * 60)
+
+
+# ---------------------------------------------------------
+# Generation
+# ---------------------------------------------------------
+
+def generate_response(
+    message,
+    history,
+    temperature,
+    top_p,
+    max_new_tokens,
+    repetition_penalty,
+):
+    if not message or not message.strip():
+        return history, ""
+
+    try:
+        # Convert Gradio history to a simple conversation
+        prompt_parts = []
+
+        for item in history:
+            if isinstance(item, dict):
+                role = item.get("role")
+                content = item.get("content", "")
+
+                if role == "user":
+                    prompt_parts.append(f"User: {content}")
+
+                elif role == "assistant":
+                    prompt_parts.append(f"Assistant: {content}")
+
+        prompt_parts.append(f"User: {message}")
+        prompt_parts.append("Assistant:")
+
+        prompt = "\n".join(prompt_parts)
+
+        inputs = tokenizer(
+            prompt,
+            return_tensors="pt",
+            truncation=True,
+            max_length=2048,
+        )
+
+        input_ids = inputs["input_ids"].to(DEVICE)
+        attention_mask = inputs["attention_mask"].to(DEVICE)
+
+        with torch.inference_mode():
+            output = model.generate(
+                input_ids=input_ids,
+                attention_mask=attention_mask,
+                max_new_tokens=int(max_new_tokens),
+                temperature=float(temperature),
+                top_p=float(top_p),
+                repetition_penalty=float(repetition_penalty),
+                do_sample=True,
+                pad_token_id=tokenizer.pad_token_id,
+                eos_token_id=tokenizer.eos_token_id,
+            )
+
+        generated_tokens = output[0][input_ids.shape[-1]:]
+
+        response = tokenizer.decode(
+            generated_tokens,
+            skip_special_tokens=True,
+        ).strip()
+
+        if not response:
+            response = "FastLLM did not generate a response."
 
+        history = history + [
+            {"role": "user", "content": message},
+            {"role": "assistant", "content": response},
+        ]
 
-@spaces.GPU(duration=30)
-def generate(prompt, max_new_tokens, temperature, top_k):
-    input_ids = tokenizer.encode(prompt, return_tensors="pt").to("cuda")
-    max_new_tokens = int(max_new_tokens)
-    temperature = float(temperature)
-    top_k = int(top_k)
+        # Release temporary tensors
+        del inputs
+        del input_ids
+        del attention_mask
+        del output
 
-    with torch.no_grad():
-        for _ in range(max_new_tokens):
-            outputs = model(input_ids)
+        if DEVICE == "cuda":
+            torch.cuda.empty_cache()
 
-            # Cast to fp32 for the sampling math — fp16 logits from an
-            # early/undertrained checkpoint can overflow fp16's range,
-            # which turns softmax output into inf/nan and crashes
-            # multinomial with a CUDA device-side assert.
-            logits = outputs["logits"][:, -1, :].float() / max(temperature, 1e-5)
+        gc.collect()
 
-            if top_k > 0:
-                v, _ = torch.topk(logits, min(top_k, logits.size(-1)))
-                logits[logits < v[:, [-1]]] = -float("Inf")
+        return history, ""
 
-            probs = torch.softmax(logits, dim=-1)
+    except Exception as e:
+        print("Generation error:")
+        print(repr(e))
 
-            # Defensive fallback: if anything still slips through as
-            # nan/inf/negative, fall back to a safe uniform distribution
-            # instead of crashing the whole generation call.
-            probs = torch.nan_to_num(probs, nan=0.0, posinf=0.0, neginf=0.0)
-            if probs.sum().item() <= 0:
-                probs = torch.ones_like(probs) / probs.size(-1)
+        error_message = f"❌ Generation error:\n\n`{str(e)}`"
 
-            next_token = torch.multinomial(probs, num_samples=1)
+        history = history + [
+            {"role": "user", "content": message},
+            {"role": "assistant", "content": error_message},
+        ]
 
-            input_ids = torch.cat([input_ids, next_token], dim=-1)
+        return history, ""
 
-            if next_token.item() == tokenizer.eos_token_id:
-                break
 
-    return tokenizer.decode(input_ids[0], skip_special_tokens=True)
+# ---------------------------------------------------------
+# Clear conversation
+# ---------------------------------------------------------
 
+def clear_chat():
+    return []
+
+
+# ---------------------------------------------------------
+# UI
+# ---------------------------------------------------------
+
+css = """
+.gradio-container {
+    max-width: 1100px !important;
+    margin: auto !important;
+}
+
+.title {
+    text-align: center;
+    margin-bottom: 4px;
+}
+
+.subtitle {
+    text-align: center;
+    opacity: 0.7;
+    margin-bottom: 20px;
+}
+
+.status {
+    text-align: center;
+    font-size: 13px;
+    opacity: 0.65;
+}
+"""
+
+with gr.Blocks(
+    css=css,
+    title="FastLLM",
+) as demo:
 
-with gr.Blocks(title="FastLLM (150M) Demo") as demo:
     gr.Markdown(
         """
-        # FastLLM (150M) — Modern Causal Language Model
-        A ~150M parameter decoder-only model (GQA, SwiGLU, RMSNorm, RoPE) from
-        [devoppro/FastLLM](https://huggingface.co/devoppro/FastLLM), still training.
-        Small model, so expect small-model quality — this is a demo, not a chatbot.
+        # ⚡ FastLLM
+        """,
+        elem_classes=["title"],
+    )
+
+    gr.Markdown(
         """
+        A 150M parameter causal language model built from scratch by **devoppro**.
+        """,
+        elem_classes=["subtitle"],
     )
+
+    chatbot = gr.Chatbot(
+        label="FastLLM",
+        height=560,
+        type="messages",
+        bubble_full_width=False,
+    )
+
     with gr.Row():
-        with gr.Column():
-            prompt = gr.Textbox(
-                label="Prompt",
-                value="Once upon a time,",
-                lines=4,
-            )
-            max_new_tokens = gr.Slider(16, 256, value=60, step=8, label="Max new tokens")
-            temperature = gr.Slider(0.1, 1.5, value=0.7, step=0.05, label="Temperature")
-            top_k = gr.Slider(0, 100, value=40, step=5, label="Top-k")
-            run_btn = gr.Button("Generate", variant="primary")
-        with gr.Column():
-            output = gr.Textbox(label="Output", lines=12)
-
-    run_btn.click(
-        fn=generate,
-        inputs=[prompt, max_new_tokens, temperature, top_k],
-        outputs=output,
+
+        message = gr.Textbox(
+            placeholder="Message FastLLM...",
+            label="",
+            scale=5,
+            lines=2,
+        )
+
+        send = gr.Button(
+            "Send",
+            variant="primary",
+            scale=1,
+        )
+
+    with gr.Row():
+
+        temperature = gr.Slider(
+            minimum=0.1,
+            maximum=2.0,
+            value=0.7,
+            step=0.05,
+            label="Temperature",
+        )
+
+        top_p = gr.Slider(
+            minimum=0.1,
+            maximum=1.0,
+            value=0.9,
+            step=0.05,
+            label="Top P",
+        )
+
+        max_tokens = gr.Slider(
+            minimum=16,
+            maximum=1024,
+            value=256,
+            step=16,
+            label="Max New Tokens",
+        )
+
+        repetition_penalty = gr.Slider(
+            minimum=1.0,
+            maximum=2.0,
+            value=1.05,
+            step=0.01,
+            label="Repetition Penalty",
+        )
+
+    with gr.Row():
+
+        clear = gr.Button(
+            "🗑️ Clear conversation"
+        )
+
+    gr.Markdown(
+        f"""
+        <div class="status">
+        Model: <b>{MODEL_ID}</b> · Device: <b>{DEVICE.upper()}</b>
+        </div>
+        """,
+        elem_classes=["status"],
+    )
+
+    # -----------------------------------------------------
+    # Events
+    # -----------------------------------------------------
+
+    send.click(
+        generate_response,
+        inputs=[
+            message,
+            chatbot,
+            temperature,
+            top_p,
+            max_tokens,
+            repetition_penalty,
+        ],
+        outputs=[
+            chatbot,
+            message,
+        ],
     )
 
-    gr.Examples(
-        examples=[
-            ["Once upon a time,", 60, 0.7, 40],
-            ["Who are you?", 60, 0.7, 40],
-            ["The most important thing about machine learning is", 60, 0.7, 40],
+    message.submit(
+        generate_response,
+        inputs=[
+            message,
+            chatbot,
+            temperature,
+            top_p,
+            max_tokens,
+            repetition_penalty,
+        ],
+        outputs=[
+            chatbot,
+            message,
         ],
-        inputs=[prompt, max_new_tokens, temperature, top_k],
     )
 
+    clear.click(
+        clear_chat,
+        inputs=[],
+        outputs=[chatbot],
+    )
+
+
+# ---------------------------------------------------------
+# Launch
+# ---------------------------------------------------------
+
 if __name__ == "__main__":
-    demo.queue().launch()
+    demo.launch(
+        server_name="0.0.0.0",
+        server_port=7860,
+    )