app.py

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