taima-brain

Taima Models

Testing, evaluation, and distillation pipeline for Taima — an AI virtual pet powered by on-device SLMs via ExecuTorch.

Repository Structure

taima-models/
├── models/                   # Downloaded .pte model files
├── results/                  # Per-model test results
├── datasets/                 # Generated fine-tuning datasets
├── ideal_responses/          # Cached ideal responses + trait scores
├── available.txt             # Registry of downloadable models
├── download_models.py        # Download .pte models from HuggingFace
├── generate_prompts.py       # Generate test prompts with pet conditions
├── test_model.py             # Run full test suite on a model
├── compare_results.py        # Aggregate and compare all results
├── generate_dataset.ipynb    # Generate training data via DeepInfra LLMs
├── distill_model.ipynb       # LoRA fine-tuning + .pte export
└── requirements.txt

Quick Start

pip install -r requirements.txt

# NOTE: model download now runs in distilation notebook
# # 1. Download a base model
# python download_models.py --list
# python download_models.py --model-id llama3.2-1b-instruct-spinquant

# 2. Generate ideal responses (run notebook or use existing)
jupyter notebook generate_dataset.ipynb

# 3. Test a model
export DEEPINFRA_API_KEY=your-key
python test_model.py --model-id llama3.2-1b-instruct-spinquant \
    --hf-model meta-llama/Llama-3.2-1B-Instruct

# 4. Compare all tested models
python compare_results.py

# 5. Fine-tune with LoRA (run notebook)
# choose between 1b and 3b models
jupyter notebook distill_model.ipynb

Test Suite

Each model is evaluated on 7 dimensions:

Overall pass threshold: 80% across all tests.

Personality Traits Evaluated

cheekiness, shyness, gentleness, grumpiness, desperation, happiness, sadness

Environment Variables

• DEEPINFRA_API_KEY — Required for LLM judge and dataset generation

ExecuTorch Setup

The .pte export requires ExecuTorch installed separately:

# Via optimum-executorch (recommended)
pip install optimum-executorch

# Or from source
git clone https://github.com/pytorch/executorch.git
cd executorch && python install_executorch.py

Running tests

python test_model.py \
  --model-id llama_3_2_1b_instruct_base_int8 \
  --tokenizer tokenizers/1b/tokenizer.model \
  --hf-model meta-llama/Llama-3.2-1B-Instruct

python test_model.py \
  --model-id llama_3_2_1b_instruct_lora_int8 \
  --tokenizer tokenizers/1b/tokenizer.model \
  --hf-model meta-llama/Llama-3.2-1B-Instruct

python test_model.py \
  --model-id llama_3_2_3b_instruct_base_int8 \
  --tokenizer tokenizers/3b/tokenizer.model \
  --hf-model meta-llama/Llama-3.2-3B-Instruct

python test_model.py \
  --model-id llama_3_2_3b_instruct_lora_int8 \
  --tokenizer tokenizers/3b/tokenizer.model \
  --hf-model meta-llama/Llama-3.2-3B-Instruct

python test_model.py \
  --model-id smollm2_135m_base_int8 \
  --tokenizer tokenizers/135m/tokenizer.model \
  --hf-model HuggingFaceTB/SmolLM2-135M-Instruct

python test_model.py \
  --model-id smollm2_135m_lora_int8 \
  --tokenizer tokenizers/135m/tokenizer.model \
  --hf-model HuggingFaceTB/SmolLM2-135M-Instruct