Building Offline-Friendly AI Features: When to Ship a Small Model On-Device

from transformers import AutoModelForCausalLM, AutoTokenizer, TrainingArguments
from trl import SFTTrainer
import torch
import torch.nn.functional as F
 
def distillation_loss(student_logits, teacher_logits, labels, temperature=3.0, alpha=0.5):
    """
    Combined loss: cross-entropy on labels + KL divergence from teacher.
    alpha: weight on distillation loss (vs standard CE loss)
    temperature: higher T = softer distributions = more transfer
    """
    # Standard supervised loss
    ce_loss = F.cross_entropy(
        student_logits.view(-1, student_logits.size(-1)),
        labels.view(-1),
        ignore_index=-100,
    )
    
    # Distillation loss
    soft_student = F.log_softmax(student_logits / temperature, dim=-1)
    soft_teacher = F.softmax(teacher_logits / temperature, dim=-1)
    kl_loss = F.kl_div(soft_student, soft_teacher, reduction="batchmean") * (temperature ** 2)
    
    return alpha * kl_loss + (1 - alpha) * ce_loss
 
class DistillationTrainer(SFTTrainer):
    def __init__(self, teacher_model, *args, **kwargs):
        super().__init__(*args, **kwargs)
        self.teacher_model = teacher_model.eval()
    
    def compute_loss(self, model, inputs, return_outputs=False):
        outputs = model(**inputs)
        
        with torch.no_grad():
            teacher_outputs = self.teacher_model(**inputs)
        
        loss = distillation_loss(
            student_logits=outputs.logits,
            teacher_logits=teacher_outputs.logits,
            labels=inputs["labels"],
        )
        
        return (loss, outputs) if return_outputs else loss

import coremltools as ct
from transformers import AutoModelForCausalLM
import torch
 
# Load and trace the model
model = AutoModelForCausalLM.from_pretrained("microsoft/phi-3-mini-4k-instruct")
model.eval()
 
# Export to ONNX first, then convert to CoreML
# (full export pipeline — abbreviated for space)
example_input = torch.randint(0, 32000, (1, 128))
 
mlmodel = ct.convert(
    traced_model,
    inputs=[ct.TensorType(shape=(1, ct.RangeDim(1, 512)), dtype=np.int32)],
    compute_precision=ct.precision.FLOAT16,  # ANE uses FP16
    compute_units=ct.ComputeUnit.ALL,        # use ANE + GPU + CPU
)
 
# Apply compression for on-device storage
op_config = ct.optimize.coreml.OpLinearQuantizerConfig(
    mode="linear_symmetric",
    dtype="int4",
    granularity="per_block",
    block_size=32,
)
config = ct.optimize.coreml.OptimizationConfig(global_config=op_config)
mlmodel_quantized = ct.optimize.coreml.linear_quantize_weights(mlmodel, config=config)
mlmodel_quantized.save("phi3-mini-int4.mlpackage")

// Android inference with ONNX Runtime Mobile
class OnDeviceLLM(context: Context) {
    private val session: OrtSession
    private val tokenizer: Tokenizer
    
    init {
        val sessionOptions = OrtSession.SessionOptions().apply {
            addNnapi()          // use Android Neural Networks API
            setInterOpNumThreads(2)
            setIntraOpNumThreads(4)
        }
        val modelBytes = context.assets.open("phi3-mini-q4.onnx").readBytes()
        session = OrtEnvironment.getEnvironment().createSession(modelBytes, sessionOptions)
        tokenizer = Tokenizer.fromAsset(context, "tokenizer.json")
    }
    
    fun generate(prompt: String, maxTokens: Int = 256): Flow<String> = flow {
        val inputIds = tokenizer.encode(prompt)
        var currentIds = inputIds.toLongArray()
        
        repeat(maxTokens) {
            val inputs = mapOf(
                "input_ids" to OnnxTensor.createTensor(
                    OrtEnvironment.getEnvironment(),
                    arrayOf(currentIds)
                )
            )
            val output = session.run(inputs)
            val logits = output[0].value as Array<Array<FloatArray>>
            val nextToken = logits[0].last().argmax()
            
            if (nextToken == tokenizer.eosTokenId) return@flow
            
            emit(tokenizer.decode(intArrayOf(nextToken)))
            currentIds = currentIds + nextToken
        }
    }
}

// iOS: Thermal and battery-aware inference throttling
import Foundation
 
class ThermalAwareInferenceManager {
    private var thermalObserver: NSObjectProtocol?
    
    enum InferenceMode {
        case full          // normal generation
        case throttled     // reduce batch, increase token delay
        case cloudFallback // offload to cloud, preserve battery
    }
    
    var currentMode: InferenceMode = .full
    
    init() {
        thermalObserver = NotificationCenter.default.addObserver(
            forName: ProcessInfo.thermalStateDidChangeNotification,
            object: nil,
            queue: .main
        ) { [weak self] _ in
            self?.updateMode()
        }
    }
    
    func updateMode() {
        let batteryLevel = UIDevice.current.batteryLevel  // -1 if unknown
        let thermalState = ProcessInfo.processInfo.thermalState
        
        switch (thermalState, batteryLevel) {
        case (.critical, _), (.serious, let b) where b < 0.20:
            currentMode = .cloudFallback
        case (.fair, let b) where b < 0.30, (.serious, _):
            currentMode = .throttled
        default:
            currentMode = .full
        }
        
        NotificationCenter.default.post(
            name: .inferenceModeChanged,
            object: currentMode
        )
    }
    
    func shouldUseOnDevice() -> Bool {
        currentMode != .cloudFallback
    }
}

import hashlib
import json
from datetime import datetime
from dataclasses import dataclass, asdict
 
@dataclass
class OfflineResponse:
    request_id: str
    query: str
    local_response: str
    timestamp: str
    model_version: str
    synced: bool = False
    cloud_response: str | None = None
 
class OfflineSyncManager:
    def __init__(self, local_store, cloud_client, divergence_threshold: float = 0.3):
        self.store = local_store
        self.cloud = cloud_client
        self.threshold = divergence_threshold
    
    def record_offline_response(self, query: str, response: str, model_ver: str) -> str:
        request_id = hashlib.sha256(f"{query}{datetime.utcnow().isoformat()}".encode()).hexdigest()[:16]
        record = OfflineResponse(
            request_id=request_id,
            query=query,
            local_response=response,
            timestamp=datetime.utcnow().isoformat(),
            model_version=model_ver,
        )
        self.store.save(request_id, asdict(record))
        return request_id
    
    async def sync_pending(self) -> list[dict]:
        unsynced = self.store.get_unsynced()
        diverged = []
        
        for record in unsynced:
            cloud_resp = await self.cloud.generate(record["query"])
            divergence = self._measure_divergence(record["local_response"], cloud_resp)
            
            record["cloud_response"] = cloud_resp
            record["synced"] = True
            self.store.update(record["request_id"], record)
            
            if divergence > self.threshold:
                diverged.append({
                    "request_id": record["request_id"],
                    "divergence": divergence,
                    "query_snippet": record["query"][:80],
                })
        
        return diverged  # caller decides whether to notify user
    
    def _measure_divergence(self, local: str, cloud: str) -> float:
        # Simple word-overlap similarity — replace with embedding similarity in production
        local_words = set(local.lower().split())
        cloud_words = set(cloud.lower().split())
        if not local_words or not cloud_words:
            return 1.0
        intersection = local_words & cloud_words
        union = local_words | cloud_words
        return 1.0 - (len(intersection) / len(union))

// iOS: Ensure on-device processing stays on-device
class PrivacyGuardedInference {
    private let onDeviceModel: OnDeviceLLM
    private let analyticsEnabled: Bool
    
    func infer(query: String) async -> String {
        // Never log the query itself — only aggregate metrics
        let startTime = Date()
        let result = await onDeviceModel.generate(query)
        let latency = Date().timeIntervalSince(startTime)
        
        if analyticsEnabled {
            // Only log performance metrics, never content
            Analytics.log("inference_completed", properties: [
                "latency_ms": latency * 1000,
                "token_count": result.split(separator: " ").count,
                "model_version": onDeviceModel.version,
                // NO: "query": query  ← never do this
                // NO: "response": result ← never do this
            ])
        }
        
        return result
    }
}