""" 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(""), ) # 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()