File size: 2,543 Bytes
ac4bbf5 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 | # 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"}
```
|