CLAUDE.md

# CLAUDE.md

This file provides guidance to Claude Code (claude.ai/code) when working with code in this repository.

## Project Overview

Python toolkit that converts Google's TranslateGemma 4B IT model (HuggingFace) into on-device inference bundles for Android. Google's official TFLite files only support WebGPU — this project produces CPU/XNNPACK-compatible `.litertlm` (LiteRT-LM) and `.task` (MediaPipe) files with proper KV-cache prefill/decode signatures.

The primary Hugging Face app artifact is the multimodal image+text LiteRT-LM bundle:
`artifacts/int4-multimodal/translategemma-4b-it-int4-multimodal.litertlm` (~2.76 GB). It supports structured text translation and image translation through LiteRT-LM vision, and is intended for PrivateAITranslate. If this file is absent from the Hugging Face repo, the app download URL for the multimodal model will 404.

## Common Commands

### Single quantization conversion (produces `.task`)
```bash
source /home/ubuntu/conv-venv/bin/activate

python convert_translategemma_android.py \
  --model-dir ./translategemma-4b-it \
  --tflite-dir ./tflite_output/dynamic_int8 \
  --output-dir ./output \
  --task-file ./output/translategemma-4b-it-native-dynamic_int8.task \
  --quantize dynamic_int8 \
  --prefill-seq-len 1024 --kv-cache-max-len 1024 --allow-no-token
```

Valid `--quantize` values: `none`, `dynamic_int8`, `int8`, `int4`, `float16`
Aliases accepted: `fp16`, `f16`, `i8`, `q8`, `i4`, `q4`, `fp32`

### Bundle a TFLite into `.litertlm` (recommended for Google AI Edge Gallery)
```bash
# Requires /tmp/litert-lm-pkg — see bundle_litertlm.py setup block if missing
python bundle_litertlm.py \
  --tflite ./tflite_output/dynamic_int8/*.tflite \
  --tokenizer ./translategemma-4b-it/tokenizer.model \
  --output ./output/translategemma-4b-it-native-dynamic_int8.litertlm \
  --quant dynamic_int8
```

### Bundle-only from existing TFLite (skip conversion)
```bash
python convert_translategemma_android.py \
  --bundle-only \
  --existing-tflite ./tflite_output/none/translategemma-4b-it-generic-none.tflite \
  --quantize int8 \
  --tflite-dir ./tflite_output/int8 \
  --output-dir ./output \
  --task-file ./output/translategemma-4b-it-int8.task
```

### Batch multi-quant build + HF upload
```bash
python multi_quant_build_upload.py \
  --model-dir ./translategemma-4b-it \
  --quants "int4,int8,dynamic_int8" \
  --repo-id barakplasma/translategemma-4b-it-android-task-quantized \
  --no-upload   # remove to upload
```

## Architecture

### Three-script design

**`convert_translategemma_android.py`** — single conversion run, three strategies in sequence:

1. **Strategy 1** (`strategy1_litert_native`) — preferred; uses `litert-torch` with `build_translategemma_4b()`, a custom builder for the 4B architecture. Produces proper KV-cache TFLite with prefill/decode signatures. Quantization is applied natively by the converter via `QUANT_MAP`:
   - `"int4"` → `"dynamic_int4_block128"` (blockwise INT4, ~2 GB)
   - `"int8"` → `"weight_only_int8"` (~4 GB)
   - `"dynamic_int8"` → `"dynamic_int8"` (~4 GB)
   - `"float16"` → `"fp16"` (~8 GB)

2. **Strategy 2** (`strategy2_generic`) — fallback; wraps the HF model in `LogitsOnlyWrapper` and exports via `ai_edge_torch.convert()`. Always exports float32. **Important**: outputs a flat `input_ids → logits` TFLite with NO KV cache — NOT compatible with MediaPipe LLM inference.

3. **Strategy 3** (`strategy3_post_tflite_quantize`) — runs only when Strategy 2 was used (never after Strategy 1); applies post-hoc weight quantization to the TFLite flatbuffer via `ai_edge_quantizer`. Does NOT add KV cache but reduces file size.

**`bundle_litertlm.py`** — takes a Strategy 1 TFLite + SentencePiece tokenizer and packages them into `.litertlm` format with `LlmMetadata` proto (Gemma3 model type, embedded Jinja chat template, BOS/EOS stop tokens, 2K max tokens). Requires `/tmp/litert-lm-pkg/` with compiled FlatBuffers and proto Python bindings (see script header for setup).

**`multi_quant_build_upload.py`** — orchestrator; invokes `convert_translategemma_android.py` as a subprocess per quant level, handles timeouts/signals, writes `output/quantization_summary.json` and `output/README.md`, uploads artifacts to HuggingFace.

### Key function: `build_translategemma_4b()`

Critical — without this, `litert-torch 0.8.0` falls back to wrong-architecture builders (1B/270m). Hardcodes the correct config:
- 34 layers, embedding_dim=2560, 8 heads, 4 KV heads, head_dim=256, intermediate=10240
- Sliding window 1024; global attention at layers where `(idx+1) % 6 == 0` (indices 5,11,17,23,29)
- RMS norm with `zero_centered=True`, per-head QK normalization (`q_norm`, `k_norm`)
- Custom loader strips `language_model.` prefix from TranslateGemma's multimodal safetensors keys (standard Gemma3 safetensors don't have this prefix)

### Output formats

| Format | Runtime | Notes |
|--------|---------|-------|
| `.litertlm` | LiteRT-LM / Google AI Edge Gallery | Recommended; embeds Jinja prompt template and LlmMetadata |
| `.task` | MediaPipe GenAI | Legacy; no embedded template — user must manually add `<start_of_turn>` tokens |

### Prompt format for on-device inference

TranslateGemma requires this exact format (trained with it):
```
<bos><start_of_turn>user
You are a professional English (en) to Spanish (es) translator...
Produce only the Spanish translation...Please translate:


{text}<end_of_turn>
<start_of_turn>model
```

In Google AI Edge Gallery **Prompt Lab** mode, paste this as the System Prompt with `{{input}}` as the placeholder. `.litertlm` files embed a simplified Jinja template for AI Chat mode.

### Runtime notes

- `conv-venv/` — virtualenv with all deps (`litert-torch==0.8.0`, `mediapipe`, `ai_edge_torch`, `transformers`)
- `/tmp/litert-lm-pkg/` — manually assembled package from cloned LiteRT-LM repo with compiled FlatBuffer (`flatc -p --gen-onefile`) and proto (`protoc`) Python bindings; required by `bundle_litertlm.py` at runtime; NOT persistent across reboots
- `/tmp/litert-lm/` — cloned `google-ai-edge/LiteRT-LM` repo (schema source for rebuilding the package)
- Multimodal conversion observed peak RSS around 73.9 GiB; require 96 GB minimum RAM, with 128 GB preferred. Ideal builder is a Vast.ai PyTorch image or Ubuntu-based NVIDIA/PyTorch image rather than a bare CUDA runtime image.
- CPU real-device validation has passed on Pixel 10 and Galaxy S22 / S22 Ultra-class targets. GPU initialization currently fails during validation and should be considered experimental and not validated.
- `translategemma-4b-it/tokenizer.model` is the SentencePiece binary used by both `.task` and `.litertlm` bundlers; `ensure_tokenizer_model()` auto-converts from `tokenizer.json` if missing
- HuggingFace repo: `barakplasma/translategemma-4b-it-android-task-quantized`; upload token in `HF_TOKEN` env var