Added Gradio app for HF space

app_smol.py

@@ -0,0 +1,223 @@
+"""
+Gradio app for SmolLM2-135M inference with streaming output.
+Uses Lightning checkpoint saved from training.
+"""
+
+import sys
+from pathlib import Path
+from typing import List, Optional
+
+import gradio as gr
+import torch
+from transformers import AutoConfig, AutoTokenizer
+
+from model import SmolConfig, SmolLM2
+from train import SmolLM2Module
+
+# Device setup
+DEVICE = "cpu"
+if torch.cuda.is_available():
+    DEVICE = "cuda"
+elif hasattr(torch.backends, "mps") and torch.backends.mps.is_available():
+    DEVICE = "mps"
+
+# Globals
+model: Optional[SmolLM2] = None
+tokenizer = None
+
+# Allow SmolConfig to be deserialized from Lightning checkpoints when torch.load
+try:
+    torch.serialization.add_safe_globals([SmolConfig])  # type: ignore[attr-defined]
+except Exception:
+    pass
+
+
+def load_model_checkpoint(checkpoint_path: str = "checkpoints/smollm2-final-step-05000.ckpt"):
+    """Load Lightning checkpoint and return status string."""
+    global model, tokenizer
+
+    ckpt = Path(checkpoint_path)
+    if not ckpt.exists():
+        return f"❌ Checkpoint not found: {ckpt}"
+
+    try:
+        hf_cfg = AutoConfig.from_pretrained("HuggingFaceTB/SmolLM2-135M")
+        config = SmolConfig.from_hf(hf_cfg)
+        tokenizer = AutoTokenizer.from_pretrained("HuggingFaceTB/SmolLM2-135M")
+        if tokenizer.pad_token is None:
+            tokenizer.pad_token = tokenizer.eos_token
+
+        module = SmolLM2Module.load_from_checkpoint(
+            str(ckpt),
+            config=config,
+            tokenizer=tokenizer,
+            map_location=DEVICE,
+            strict=False,
+        )
+        module.eval()
+        model = module.model.to(DEVICE).eval()
+        return f"✅ Model loaded from {ckpt} on {DEVICE}"
+    except Exception as e:  # pragma: no cover - interactive
+        model = None
+        return f"❌ Error loading model: {e}"
+
+
+def stream_generate(
+    prompt: str,
+    max_new_tokens: int,
+    temperature: float,
+    top_k: int,
+    top_p: float,
+):
+    """Generator that yields only the generated text (without prompt)."""
+    global model, tokenizer
+    if model is None or tokenizer is None:
+        yield "⚠️ Load the model first (click Reload Model)."
+        return
+
+    if not prompt or not prompt.strip():
+        yield "⚠️ Please enter a prompt."
+        return
+
+    # Tokenize prompt
+    inputs = tokenizer(prompt, return_tensors="pt", add_special_tokens=False)
+    input_ids = inputs["input_ids"].to(DEVICE)
+
+    # Guard against context overflow
+    if input_ids.shape[1] >= model.config.max_position_embeddings:
+        yield f"⚠️ Prompt too long ({input_ids.shape[1]} tokens). Max is {model.config.max_position_embeddings}."
+        return
+
+    generated = input_ids
+    past_key_values: Optional[List] = None
+    prompt_length = input_ids.shape[1]
+
+    with torch.no_grad():
+        for _ in range(max_new_tokens):
+            if past_key_values is None:
+                current_input = generated
+            else:
+                current_input = generated[:, -1:]
+
+            logits, past_key_values = model(
+                current_input,
+                past_key_values=past_key_values,
+                use_cache=True,
+            )
+
+            next_token_logits = logits[:, -1, :] / max(temperature, 1e-6)
+
+            # top-k
+            if top_k > 0:
+                values, _ = torch.topk(next_token_logits, top_k)
+                min_keep = values[:, -1].unsqueeze(-1)
+                next_token_logits = torch.where(
+                    next_token_logits < min_keep,
+                    torch.full_like(next_token_logits, float("-inf")),
+                    next_token_logits,
+                )
+
+            # top-p
+            if top_p < 1.0:
+                sorted_logits, sorted_indices = torch.sort(next_token_logits, descending=True)
+                probs = torch.softmax(sorted_logits, dim=-1)
+                cumulative = torch.cumsum(probs, dim=-1)
+                sorted_mask = cumulative > top_p
+                sorted_mask[..., 1:] = sorted_mask[..., :-1].clone()
+                sorted_mask[..., 0] = 0
+                mask = sorted_mask.scatter(1, sorted_indices, sorted_mask)
+                next_token_logits = torch.where(mask, torch.full_like(next_token_logits, float("-inf")), next_token_logits)
+
+            probs = torch.softmax(next_token_logits, dim=-1)
+            next_token = torch.multinomial(probs, num_samples=1)
+
+            generated = torch.cat([generated, next_token], dim=1)
+            # Decode only the generated part (skip the prompt)
+            generated_text = tokenizer.decode(generated[0][prompt_length:], skip_special_tokens=True)
+            yield generated_text
+
+
+# Initial load
+INITIAL_STATUS = load_model_checkpoint()
+
+
+def chat_stream(message, history, max_tokens, temperature, top_k, top_p):
+    """Gradio wrapper for streaming chat."""
+    if history is None:
+        history = []
+
+    # Convert history from tuple format to dict format if needed
+    if history and isinstance(history[0], (list, tuple)):
+        # Convert from tuple format [(user, assistant), ...] to dict format
+        new_history = []
+        for h in history:
+            if isinstance(h, (list, tuple)) and len(h) >= 2:
+                if h[0]:  # User message
+                    new_history.append({"role": "user", "content": str(h[0])})
+                if h[1]:  # Assistant message
+                    new_history.append({"role": "assistant", "content": str(h[1])})
+        history = new_history
+
+    # Append user message
+    user_msg = (message or "").strip()
+    if not user_msg:
+        yield history
+        return
+    
+    history.append({"role": "user", "content": user_msg})
+    history.append({"role": "assistant", "content": ""})
+
+    stream = stream_generate(user_msg, max_tokens, temperature, top_k, top_p)
+    for partial in stream:
+        # Update the last assistant message with generated text
+        if partial:
+            history[-1] = {"role": "assistant", "content": str(partial)}
+        yield history
+
+
+def clear_chat():
+    return "", []
+
+
+with gr.Blocks(title="SmolLM2-135M Text Generator") as demo:
+    gr.Markdown(
+        """
+        # 🤖 SmolLM2-135M Text Generator
+
+        Generate text with your trained SmolLM2-135M checkpoint (streaming output).
+        """
+    )
+
+    with gr.Row():
+        with gr.Column(scale=1):
+            gr.Markdown("### Model Status")
+            status_text = gr.Textbox(value=INITIAL_STATUS, label="Status", interactive=False, lines=2)
+            load_btn = gr.Button("🔄 Reload Model", variant="secondary")
+            ckpt_input = gr.Textbox(
+                value="checkpoints/smollm2-step=05000-train_loss=0.0918.ckpt",
+                label="Checkpoint path",
+                interactive=True,
+            )
+            load_btn.click(fn=lambda p: load_model_checkpoint(p), inputs=ckpt_input, outputs=status_text)
+
+            gr.Markdown("### Generation Parameters")
+            max_tokens = gr.Slider(10, 500, value=100, step=10, label="Max Tokens")
+            temperature = gr.Slider(0.1, 2.0, value=0.8, step=0.1, label="Temperature")
+            top_k = gr.Slider(0, 100, value=50, step=5, label="Top-K")
+            top_p = gr.Slider(0.1, 1.0, value=1.0, step=0.05, label="Top-P")
+
+        with gr.Column(scale=2):
+            gr.Markdown("### 💬 Chat Interface")
+            chatbot = gr.Chatbot(label="Conversation", height=500)
+            with gr.Row():
+                msg = gr.Textbox(label="Your Message", placeholder="Type your prompt here...", scale=4, lines=2)
+                submit_btn = gr.Button("Send ➤", variant="primary", scale=1)
+            clear_btn = gr.Button("🗑️ Clear Chat", variant="stop")
+
+    msg.submit(fn=chat_stream, inputs=[msg, chatbot, max_tokens, temperature, top_k, top_p], outputs=chatbot)
+    submit_btn.click(fn=chat_stream, inputs=[msg, chatbot, max_tokens, temperature, top_k, top_p], outputs=chatbot).then(fn=lambda: "", outputs=msg)
+    clear_btn.click(fn=clear_chat, outputs=[msg, chatbot])
+
+
+if __name__ == "__main__":
+    demo.queue().launch(share=False, server_name="0.0.0.0", server_port=7860)

logs/tensorboard/version_3/hparams.yaml

@@ -0,0 +1,5 @@
+block_size: 512
+peak_lr: 0.0005
+predict_every: 500
+total_steps: 5000
+warmup_steps: 1000

logs/training_20251209_154910.log

@@ -33,3 +33,48 @@ First Citizen:
 None,
 2025-12-09 15:59:47,488 - INFO - ================================================================================
 
+2025-12-09 16:10:06,586 - INFO - Step 2500 | train_loss=0.9911
+2025-12-09 16:10:08,637 - INFO - 
+================================================================================
+2025-12-09 16:10:08,637 - INFO - Step 2500 - Generated text:
+2025-12-09 16:10:08,637 - INFO - First Citizen:
+He said he: youCLARENCE:
+He hath nopt to die among this case,
+Yet to flatter, shield your wit would not have not right.
+
+LADY ANNE:
+It is it so.
+2025-12-09 16:10:08,637 - INFO - ================================================================================
+
+2025-12-09 16:20:02,546 - INFO - Step 3000 | train_loss=0.6307
+2025-12-09 16:20:04,468 - INFO - 
+================================================================================
+2025-12-09 16:20:04,468 - INFO - Step 3000 - Generated text:
+2025-12-09 16:20:04,468 - INFO - First Citizen:
+Come, let us go in our delay: if
+you guard guard guard Corioli, your rash a
+more in yourple; even your need, the queen,
+Your wives,
+Your loving, bosom, kill into his
+2025-12-09 16:20:04,468 - INFO - ================================================================================
+
+2025-12-09 16:30:01,255 - INFO - Step 3500 | train_loss=0.1352
+2025-12-09 16:30:03,305 - INFO - 
+================================================================================
+2025-12-09 16:30:03,305 - INFO - Step 3500 - Generated text:
+2025-12-09 16:30:03,305 - INFO - First Citizen:
+Nor I.
+
+CORIOLANUS:
+Not now, if it be your will be here.
+
+MENENIUS:
+I tell thee, fellow,
+If thou dost love to see thee,
+
+2025-12-09 16:30:03,305 - INFO - ================================================================================
+
+2025-12-09 16:30:18,743 - INFO - Final checkpoint saved: checkpoints/smollm2-final-step-03500.ckpt
+2025-12-09 16:31:03,806 - INFO - Training completed!
+2025-12-09 16:31:03,807 - INFO - Best checkpoint: /Users/qureshsu/Learning/TSAI/ERAV4/session13/smolLM-135/checkpoints/smollm2-step=03500-train_loss=0.1352.ckpt
+2025-12-09 16:31:03,807 - INFO - Last checkpoint: /Users/qureshsu/Learning/TSAI/ERAV4/session13/smolLM-135/checkpoints/last.ckpt

logs/training_20251209_163515.log

@@ -0,0 +1,71 @@
+2025-12-09 16:35:15,206 - INFO - Logging to: logs/training_20251209_163515.log
+2025-12-09 16:35:15,206 - INFO - Loading tokenizer...
+2025-12-09 16:35:16,040 - INFO - Loading model config...
+2025-12-09 16:35:16,277 - INFO - Loading dataset from: /Users/qureshsu/Learning/TSAI/ERAV4/session13/data/input.txt
+2025-12-09 16:35:16,738 - INFO - Initializing model...
+2025-12-09 16:35:17,466 - INFO - Starting training...
+2025-12-09 16:35:17,466 - INFO - Resuming from checkpoint: checkpoints/smollm2-step=03500-train_loss=0.1352.ckpt
+2025-12-09 16:35:35,153 - INFO - 
+================================================================================
+2025-12-09 16:35:35,153 - INFO - MODEL SUMMARY
+2025-12-09 16:35:35,153 - INFO - ================================================================================
+2025-12-09 16:35:35,153 - INFO - Model: SmolLM2-135M
+2025-12-09 16:35:35,153 - INFO - Total parameters: 134,515,008
+2025-12-09 16:35:35,153 - INFO - Trainable parameters: 134,515,008
+2025-12-09 16:35:35,153 - INFO - Block size: 512
+2025-12-09 16:35:35,153 - INFO - Warmup steps: 1000
+2025-12-09 16:35:35,153 - INFO - Peak learning rate: 0.0005
+2025-12-09 16:35:35,153 - INFO - Total training steps: 5000
+2025-12-09 16:35:35,153 - INFO - Predict every: 500 steps
+2025-12-09 16:35:35,153 - INFO - ================================================================================
+
+2025-12-09 16:46:13,641 - INFO - Step 4000 | train_loss=0.5093
+2025-12-09 16:46:15,889 - INFO - 
+================================================================================
+2025-12-09 16:46:15,890 - INFO - Step 4000 - Generated text:
+2025-12-09 16:46:15,890 - INFO - First Citizen:
+What a strange news, what he hath done famously
+All slain and g indeed.
+
+KING HENRY VI:
+Hadst thou been kill'd, I would not sh wrong;
+And by that you are, some thou
+2025-12-09 16:46:15,890 - INFO - ================================================================================
+
+2025-12-09 16:56:44,602 - INFO - Step 4500 | train_loss=0.5634
+2025-12-09 16:56:46,770 - INFO - 
+================================================================================
+2025-12-09 16:56:46,770 - INFO - Step 4500 - Generated text:
+2025-12-09 16:56:46,770 - INFO - First Citizen:
+'Tis a nupt in a sword's make him my
+First Citizen:
+Therefore.
+
+First Citizen:
+Is there no hope?
+
+Third Citizen:
+And ta'en! Suffolk, we shall bring all
+
+2025-12-09 16:56:46,770 - INFO - ================================================================================
+
+2025-12-09 17:07:03,502 - INFO - Step 5000 | train_loss=0.0918
+2025-12-09 17:07:06,185 - INFO - 
+================================================================================
+2025-12-09 17:07:06,186 - INFO - Step 5000 - Generated text:
+2025-12-09 17:07:06,186 - INFO - First Citizen:
+You must think of it?
+
+Pedant:
+Ay, I have
+AUTOLYCUS:
+Pray you, who came George to 't last once.
+
+AUTOLYCUS:
+I know
+2025-12-09 17:07:06,186 - INFO - ================================================================================
+
+2025-12-09 17:07:18,753 - INFO - Final checkpoint saved: checkpoints/smollm2-final-step-05000.ckpt
+2025-12-09 17:07:49,059 - INFO - Training completed!
+2025-12-09 17:07:49,060 - INFO - Best checkpoint: /Users/qureshsu/Learning/TSAI/ERAV4/session13/smolLM-135/checkpoints/smollm2-step=05000-train_loss=0.0918.ckpt
+2025-12-09 17:07:49,060 - INFO - Last checkpoint: /Users/qureshsu/Learning/TSAI/ERAV4/session13/smolLM-135/checkpoints/last.ckpt

pyproject.toml

@@ -14,4 +14,5 @@ dependencies = [
     "torchvision>=0.24.1",
     "tqdm>=4.67.1",
     "transformers>=4.57.3",
+    "gradio>=4.44.0",
 ]

train.py

@@ -234,9 +234,9 @@ def main():
     block_size = 512
     batch_size = 4
     num_workers = 8
-    max_steps = 3500
+    max_steps = 5000
     predict_every = 500
-    resume_from_checkpoint = "checkpoints/smollm2-step=01500-train_loss=3.6240.ckpt"  # Set to checkpoint path to resume, or None for fresh training
+    resume_from_checkpoint = "checkpoints/smollm2-step=03500-train_loss=0.1352.ckpt"  # Set to checkpoint path to resume, or None for fresh training
     
     # Training hyperparameters from paper
     warmup_steps = 1000