Upload app.py with huggingface_hub

app.py

@@ -0,0 +1,259 @@
+"""
+Tiny-LLM CLI SFT Demo - Generate Shell Commands from Natural Language
+
+This model was fine-tuned to translate natural language instructions to CLI commands.
+"""
+
+import gradio as gr
+import torch
+from huggingface_hub import hf_hub_download
+from model import TinyLLM, MODEL_CONFIG
+
+# Model configuration
+MODEL_ID = "jonmabe/tiny-llm-cli-sft"
+MODEL_FILENAME = "best_model.pt"
+
+# Load tokenizer
+try:
+    from tokenizers import Tokenizer
+    tokenizer_path = hf_hub_download(repo_id=MODEL_ID, filename="tokenizer.json")
+    tokenizer = Tokenizer.from_file(tokenizer_path)
+    print("Loaded tokenizer from model repo")
+except Exception as e:
+    print(f"Could not load tokenizer: {e}")
+    tokenizer = None
+
+# Load model
+print("Downloading model...")
+model_path = hf_hub_download(repo_id=MODEL_ID, filename=MODEL_FILENAME)
+print(f"Model downloaded to {model_path}")
+
+print("Loading model...")
+checkpoint = torch.load(model_path, map_location="cpu", weights_only=False)
+
+# Get config from checkpoint if available
+if "config" in checkpoint and isinstance(checkpoint["config"], dict):
+    config = checkpoint["config"]
+    if "model" in config:
+        config = config["model"]
+else:
+    config = MODEL_CONFIG
+
+# Initialize model
+model = TinyLLM(config)
+
+# Load weights
+if "model_state_dict" in checkpoint:
+    state_dict = checkpoint["model_state_dict"]
+else:
+    state_dict = checkpoint
+
+missing, unexpected = model.load_state_dict(state_dict, strict=False)
+if missing:
+    print(f"Warning: Missing keys: {missing[:5]}...")
+if unexpected:
+    print(f"Warning: Unexpected keys: {unexpected[:5]}...")
+
+# Move to device
+device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+model = model.to(device)
+model.eval()
+
+total_params = sum(p.numel() for p in model.parameters())
+print(f"Model loaded on {device} with {total_params:,} parameters")
+
+
+def clean_bpe_output(text: str) -> str:
+    """Clean BPE artifacts from tokenizer output."""
+    # Replace BPE space marker with actual space
+    text = text.replace("Ġ", " ")
+    # Replace BPE newline marker with actual newline
+    text = text.replace("Ċ", "\n")
+    # Clean up extra spaces
+    text = " ".join(text.split())
+    return text.strip()
+
+
+def generate_command(
+    instruction: str,
+    max_tokens: int = 50,
+    temperature: float = 0.7,
+    top_p: float = 0.9,
+    top_k: int = 50,
+) -> str:
+    """Generate a CLI command from an instruction."""
+    
+    if not instruction.strip():
+        return "Please enter an instruction."
+    
+    if tokenizer is None:
+        return "Tokenizer not available."
+    
+    # Format prompt
+    prompt = f"Instruction: {instruction}\nCommand:"
+    
+    # Tokenize
+    encoded = tokenizer.encode(prompt)
+    input_ids = torch.tensor([encoded.ids], dtype=torch.long).to(device)
+    input_len = input_ids.shape[1]
+    
+    # Generate
+    with torch.no_grad():
+        output_ids = model.generate(
+            input_ids,
+            max_new_tokens=max_tokens,
+            temperature=temperature,
+            top_p=top_p,
+            top_k=top_k,
+            eos_token_id=tokenizer.token_to_id("</s>"),
+        )
+    
+    # Decode only the generated tokens
+    generated_ids = output_ids[0, input_len:].tolist()
+    raw_output = tokenizer.decode(generated_ids)
+    
+    # Clean BPE artifacts
+    command = clean_bpe_output(raw_output)
+    
+    # Extract just the command (first line, stop at newline)
+    command = command.split("\n")[0].strip()
+    
+    return command
+
+
+# Example instructions
+EXAMPLES = [
+    ["List all files in the current directory"],
+    ["Find all Python files"],
+    ["Show disk usage"],
+    ["Create a new folder called test"],
+    ["Search for 'error' in log files"],
+    ["Show the last 10 lines of a file"],
+    ["Count lines in a file"],
+    ["Copy files to another directory"],
+    ["Show running processes"],
+    ["Check available disk space"],
+]
+
+
+# Create Gradio interface
+with gr.Blocks(title="CLI Command Generator") as demo:
+    gr.Markdown("""
+    # 🖥️ CLI Command Generator
+    
+    Translate natural language instructions to shell commands using a **54M parameter** language model.
+    
+    ⚠️ **Note**: This is an early-stage SFT model. Outputs may be incomplete or incorrect.
+    
+    ### How to Use
+    1. Enter a natural language instruction
+    2. Click "Generate" or press Enter
+    3. The model will suggest a shell command
+    """)
+    
+    with gr.Row():
+        with gr.Column(scale=2):
+            instruction_input = gr.Textbox(
+                label="Instruction",
+                placeholder="Describe what you want to do...",
+                lines=2,
+                value="List all files in the current directory"
+            )
+            
+            with gr.Row():
+                with gr.Column():
+                    max_tokens = gr.Slider(
+                        minimum=10,
+                        maximum=100,
+                        value=50,
+                        step=5,
+                        label="Max Tokens",
+                    )
+                    temperature = gr.Slider(
+                        minimum=0.1,
+                        maximum=1.5,
+                        value=0.7,
+                        step=0.1,
+                        label="Temperature",
+                        info="Higher = more creative"
+                    )
+                
+                with gr.Column():
+                    top_p = gr.Slider(
+                        minimum=0.1,
+                        maximum=1.0,
+                        value=0.9,
+                        step=0.05,
+                        label="Top-p",
+                    )
+                    top_k = gr.Slider(
+                        minimum=1,
+                        maximum=100,
+                        value=50,
+                        step=5,
+                        label="Top-k",
+                    )
+            
+            generate_btn = gr.Button("⚡ Generate Command", variant="primary", size="lg")
+        
+        with gr.Column(scale=2):
+            output_command = gr.Textbox(
+                label="Generated Command",
+                lines=3,
+                interactive=False,
+            )
+            
+            gr.Markdown("""
+            ### Common Commands Reference
+            - `ls` - list files
+            - `find` - search for files
+            - `grep` - search in files
+            - `df` - disk usage
+            - `du` - directory size
+            - `tar` - archive files
+            - `scp` - copy over SSH
+            """)
+    
+    gr.Markdown("### 📝 Example Instructions")
+    gr.Examples(
+        examples=EXAMPLES,
+        inputs=instruction_input,
+    )
+    
+    # Event handlers
+    generate_btn.click(
+        fn=generate_command,
+        inputs=[instruction_input, max_tokens, temperature, top_p, top_k],
+        outputs=output_command,
+    )
+    
+    instruction_input.submit(
+        fn=generate_command,
+        inputs=[instruction_input, max_tokens, temperature, top_p, top_k],
+        outputs=output_command,
+    )
+    
+    gr.Markdown("""
+    ---
+    ### About This Model
+    
+    **Model**: [jonmabe/tiny-llm-cli-sft](https://huggingface.co/jonmabe/tiny-llm-cli-sft)
+    
+    This is a Supervised Fine-Tuned (SFT) version of [tiny-llm-54m](https://huggingface.co/jonmabe/tiny-llm-54m),
+    trained on ~13,000 natural language → CLI command pairs.
+    
+    #### Known Limitations
+    - 🔬 **Experimental**: Outputs may be incomplete or incorrect
+    - 📊 **Small model**: 54M parameters limits capability
+    - 🔧 **Needs improvement**: More training data and steps needed
+    
+    #### Training Details
+    - **Steps**: 2,000
+    - **Best Val Loss**: 1.2456
+    - **Data**: Geddy's NL2Bash + NL2Bash benchmark + synthetic
+    - **Hardware**: RTX 5090, ~9 minutes
+    """)
+
+
+if __name__ == "__main__":
+    demo.launch()