# API & Inference Usage

This guide covers how to load the MiniLM 1.58-bit base model and dynamically snap on custom LoRAs for inference.

## Python Inference (PyTorch)

Because MiniLM uses custom ternary BitLinear layers, it cannot be loaded via the standard `transformers` AutoModel pipeline. You must use the provided `model.py` and `lora.py` scripts.

### 1. Loading the Base Model
```python
import torch
from transformers import AutoTokenizer
from model import BitGPT

device = torch.device('mps' if torch.backends.mps.is_available() else 'cpu')
tokenizer = AutoTokenizer.from_pretrained("HuggingFaceTB/SmolLM-135M-Instruct")

# Initialize the 12-Layer Tied Architecture
model = BitGPT(vocab_size=len(tokenizer), embed_dim=256, num_layers=12, num_heads=4, tie_weights=True).to(device)

# Load the frozen 1.58-bit Base Weights
model.load_state_dict(torch.load("minilm_base.pt", map_location=device))
model.eval()
```

### 2. Injecting a "Side-Car" LoRA
If you want to run a specific task (like Smart Home JSON extraction), you must wrap the Linear layers with the custom `BitLoraLinear` adapter.

```python
from lora import inject_lora

# Wrap the model's layers with LoRA adapters
model = inject_lora(model, r=8, lora_alpha=16).to(device)

# Snap on the custom 1MB weights (strict=False ensures we only overwrite the new LoRA parameters)
model.load_state_dict(torch.load("lora_smarthome.pt", map_location=device), strict=False)
model.eval()
```

### 3. Generation Loop
To generate text, format your prompt using `ChatML` standard tags:

```python
prompt = "Uh, it's freezing in here, can you turn up the heat in the living room?"
chatml_text = f"<|im_start|>user\n{prompt}<|im_end|>\n<|im_start|>assistant\n"
input_ids = tokenizer.encode(chatml_text, return_tensors="pt").to(device)

max_new_tokens = 60
with torch.no_grad():
    for _ in range(max_new_tokens):
        logits = model(input_ids)
        next_token_logits = logits[:, -1, :]
        
        # Greedy decoding
        next_token = torch.argmax(next_token_logits, dim=-1, keepdim=True)
        input_ids = torch.cat([input_ids, next_token], dim=-1)
        
        # Stop condition (2 is im_end in ChatML)
        if next_token.item() == tokenizer.eos_token_id or next_token.item() == 2:
            break
            
output_text = tokenizer.decode(input_ids[0])
final_output = output_text.split("<|im_start|>assistant\n")[-1].replace("<|im_end|>", "").strip()

print(final_output)
# Output: {"device": "thermostat", "action": "increase_temp", "room": "living_room"}
```