gemma-4-E2B-it-AgroRAG-LiteRT

🌾 Guia do Formato .litertlm & Abordagens de Modelos de Borda no AgroRAG

Este documento detalha o funcionamento técnico do formato de modelo .litertlm, o pipeline de conversão utilizado neste projeto, e compara as três principais abordagens avaliadas para rodar Modelos de Linguagem de Pequeno Porte (SLMs) de forma 100% offline em dispositivos móveis no projeto AgroRAG.

---

🗺️ 1. O que é o formato .litertlm?

O formato .litertlm é uma especificação desenvolvida pela comunidade do TensorFlow Lite (LiteRT) da Google, voltada especificamente para empacotar e executar Grandes Modelos de Linguagem (LLMs/SLMs) de maneira otimizada em dispositivos de borda (Edge AI).

Diferente de formatos tradicionais de tensores puros (como .safetensors), o .litertlm funciona como um contêiner unificado (bundle). Ele agrupa:

1. Metadata do LLM (LlmMetadataProto): Configurações de hiperparâmetros, tokens especiais, e regras de parada.
2. Tokenizador: Configuração e tabelas de tokenização (ex: HF_Tokenizer_Zlib).
3. Grafos de Computação (.tflite):
   • tf_lite_embedder: Módulo de projeção inicial dos embeddings de entrada.
   • tf_lite_per_layer_embedder: Projeção de embeddings em camadas internas.
   • tf_lite_prefill_decode: O grafo principal do modelo que computa as etapas de prefill (processamento do prompt) e decode (geração de tokens sequenciais).
4. Módulos de Áudio e Visão (opcionais): Componentes para suporte multimodal.
5. Drafter (mtp_drafter): Modelo menor opcional para decodificação especulativa (Speculative Decoding).

Benefícios do .litertlm

• Compilação de Operações Complexas: Compila grafos matemáticos pesados diretamente para instruções nativas ARM e alocação dinâmica de buffers na GPU/NPU.
• Aceleração de Hardware Nativa: Comunicação de baixo nível através de drivers OpenCL, Vulkan ou a NNAPI do Android, permitindo velocidades de geração elevadas (>18 tokens/segundo) com consumo moderado de bateria.

---

🔄 2. O Pipeline de Fusão e Exportação .litertlm

Para utilizar o modelo Gemma 4 E2B IT adaptado ao domínio agrícola, o projeto implementou um pipeline de duas fases executado em ambientes com alta disponibilidade de memória RAM (recomenda-se no mínimo 40 GB para o Merge e 80 GB para o Export).

O pipeline é orquestrado pelo script Pipeline_gemma.py. O notebook de referência fine-tuned-gemma-4-on-mobile-lite-rt-solution.ipynb é mantido sob a pasta de treinamento estritamente como uma referência técnica externa sobre os desafios de engenharia reversa do formato.

[!NOTE]
É importante destacar que o treinamento de ajuste fino (Fine-Tuning) foi realizado previamente, gerando os adaptadores LoRA hospedados no repositório oficial do Hugging Face: ThiagoAllan/gemma-4-E2B-it-AgroRAG.
O script Pipeline_gemma.py não executa o treinamento do zero. Ele simplesmente realiza o download dos arquivos necessários para a conversão (os pesos do LoRA que residem sob a subpasta lora do repositório) e o modelo base (google/gemma-4-E2B-it), fundindo-os localmente para gerar o bundle .litertlm. Todo o treinamento do AgroRAG foi realizado exatamente com base nesta estrutura.

Visualização do Fluxo de Conversão

┌──────────────────────┐               ┌──────────────────────────────────────┐
│  Gemma 4 E2B Base    │               │ HF: ThiagoAllan/gemma-4-E2B-it-...   │
└──────────┬───────────┘               └──────────────────┬───────────────────┘
           │                                              │
     (Carrega na RAM)                             (Download Automático)
           │                                              │
           ▼                                              ▼
┌─────────────────────────────────────────────────────────────────────────────┐
│                         Fusão dos Pesos (Merge na RAM)                      │
└──────────────────────────────────────┬──────────────────────────────────────┘
                                       │
                            (save_pretrained_merged)
                                       │
                                       ▼
┌─────────────────────────────────────────────────────────────────────────────┐
│                         Modelo Fundido FP16 Local                           │
└──────────────────────────────────────┬──────────────────────────────────────┘
                                       │
                       (litert-torch export_hf via CLI)
                                       │
                                       ▼
┌─────────────────────────────────────────────────────────────────────────────┐
│                       Geração dos Sub-gráficos .tflite                      │
└──────────────────────────────────────┬──────────────────────────────────────┘
                                       │
                                       ▼
┌─────────────────────────────────────────────────────────────────────────────┐
│                             { Custom Bundler }                              │
└──────────────────────────────────────▲──────────────────────────────────────┘
                                       │
                         (Extração de seções limpas)
                                       │
┌──────────────────────────────────────┴──────────────────────────────────────┐
│                        Gemma 4 E2B Base .litertlm                           │
└─────────────────────────────────────────────────────────────────────────────┘
                                       │
           (Mescla Embedders + Prefill Re-exportado + Seções Base)
                                       │
                                       ▼
┌─────────────────────────────────────────────────────────────────────────────┐
│                            model.litertlm Final                             │
└─────────────────────────────────────────────────────────────────────────────┘

Detalhamento das Etapas

1. Fase 1: Fusão dos Pesos (Merge) O script Pipeline_gemma.py faz o download dos adaptadores a partir do repositório ThiagoAllan/gemma-4-E2B-it-AgroRAG e os acopla ao modelo base utilizando a biblioteca PEFT.

   [!IMPORTANT]
   Deve-se utilizar model.save_pretrained_merged() (como feito no notebook de referência) ou garantir que os pesos quantizados de 4 bits sejam adequadamente desquantizados para float16 na RAM antes de serem salvos, evitando que o conversor da Google quebre ou gere pesos corrompidos.

2. Fase 2: Exportação e Quantização (litert-torch) A ferramenta de CLI litert-torch (chamada via subprocesso no script Pipeline_gemma.py) rastreia o modelo local para exportá-lo como módulos TFLite. O comando utilizado internamente é:

   litert-torch export_hf      --model=./gemma_4_fusion_local      --output_dir=./gemma_4_RAG_mobile      --task=image_text_to_text      --export_vision_encoder=False      --keep_temporary_files=True      --quantization_recipe=dynamic_wi8_afp32      --externalize_embedder
   

---

⚠️ 3. Regras de Ouro e Lições Aprendidas no Desenvolvimento do .litertlm

Durante a engenharia de exportação do modelo no projeto, foram descobertas armadilhas técnicas cruciais que afetam a estabilidade do aplicativo móvel:

Regra 1: Coerência de Pipeline dos Embedders (A Divisa EvalHybrid)

• O Problema: O modelo base da Google .litertlm oficial foi gerado por compiladores internos proprietários, resultando em um embedder altamente compactado (~104 MB). Ao tentar economizar tempo de exportação e fatiar (splice) o embedder base da Google para combinar com o grafo de linguagem (prefill_decode) re-exportado pelo litert-torch, o modelo compilará com sucesso, mas sofrerá uma falha crítica de segmentação (Segmentation Fault) ou gerará tokens corrompidos no dispositivo móvel.
• A Solução: Os blocos de linguagem (tf_lite_embedder, tf_lite_per_layer_embedder e tf_lite_prefill_decode) devem vir da mesma invocação do litert-torch export_hf orquestrada em Pipeline_gemma.py. O custom bundler deve mesclar essas 3 peças geradas de uma vez, juntamente com o tokenizer do checkpoint ajustado.

Regra 2: Uso do Task type image_text_to_text (Mesmo para Text-only)

• O Problema: Rastrear o modelo usando --task text_generation gera assinaturas incompatíveis de entrada no buffer interno do interpretador TFLite do Gemma 4, fazendo o modelo emitir sequências aleatórias/ruídos.
• A Solução: Usar sempre a tarefa multimodal --task image_text_to_text combinada com --export_vision_encoder=False para pular a compilação de visão do zero (que carrega bugs de broadcast com numpy em quantizações dinâmicas). O custom bundler posteriormente enxerta a seção de visão original do modelo base, que permanece inalterada.

Regra 3: Preservação do Persona do LoRA (wi8 vs wi4)

• O Problema: A quantização agressiva de pesos em 4 bits (dynamic_wi4_afp32) reduz o tamanho do modelo para ~2.7 GB, porém remove a variação matemática fina introduzida pelo ajuste fino de domínio (LoRA). Na prática, o modelo reverte à sua identidade padrão ("Eu sou a Gemma, assistente da Google") e ignora as instruções factuais do prompt RAG.
• A Solução: Utilizar a receita dynamic_wi8_afp32 (pesos em int8 e ativações dinâmicas de float32 mapeadas em tempo de execução). Com precisão de 8 bits para os pesos, o delta do LoRA permanece intacto no espaço latente, garantindo obediência estrita ao prompt de contexto agronômico da Embrapa.

---

📊 4. Estrutura Interna de Seções do Bundle .litertlm

Um bundle final customizado e funcional (model-dyn-wi8-afp32.litertlm) possui a seguinte estrutura interna de 12 seções, cuja ordem deve ser respeitada estritamente para o correto funcionamento da máquina de estado da GPU móvel:

---

🏆 5. Comparativo das Três Abordagens do Projeto (Every Approach)

Durante o ciclo de vida do AgroRAG, três abordagens distintas foram desenhadas e implementadas para realizar inferência offline de SLMs. O quadro abaixo detalha as características e o motivo de seleção de cada uma:

gantt
    title Evolução e Abordagens de SLMs no AgroRAG
    dateFormat  YYYY-MM-DD
    section Fase Inicial
    Gemma 4 E2B (.litertlm) via LiteRT :a1, 2026-01-01, 90d
    section Fase Intermediária
    Gemma 4 E2B (GGUF) via llama.rn :a2, 2026-03-01, 60d
    section Fase Final (Produção)
    Qwen 2.5 1.5B (GGUF) via llama.rn :crit, a3, 2026-05-01, 52d

Raciocínio por trás da Abordagem Final (Qwen 2.5 1.5B GGUF)

Embora os testes conceituais com o Gemma 4 em formato .litertlm tenham apresentado velocidades impressionantes em GPUs de última geração (como nos aparelhos topo de linha), a realidade rural exige o suporte a aparelhos intermediários antigos (SoC Qualcomm Snapdragon série 7 ou inferior, com apenas 4 GB a 6 GB de memória física).

A transição técnica para o Qwen 2.5 1.5B GGUF executado via llama.rn foi consolidada na versão de produção pelos seguintes fatores:

1. Redução Drástica de Memória: O modelo ocupa apenas 1.02 GB in disk e menos de 2 GB de RAM ativa, erradicando os travamentos causados pelo coletor de processos do Android (Out-Of-Memory Killer).
2. Facilidade de Depuração: O ecossistema GGUF possui suporte consolidado de ferramentas de quantização comunitárias via Llama.cpp, sem a necessidade de remendos em scripts MLIR/TensorFlow de compilação proprietária.
3. Consistência Semântica: O ajuste fino supervisionado (SFT) de 3 épocas com Unsloth sobre o Qwen 2.5 permitiu ao modelo seguir rigidamente as instruções de "Não alucinar" (falha segura), gerando listas formatadas em bullet points de maneira mais limpa e determinística.

---

🛠️ 6. Como Inspecionar seu Bundle .litertlm

Para inspecionar um arquivo .litertlm local e validar se a estrutura de metadados e os tamanhos das seções estão coerentes (evitando o problema de embedder de tamanho incompatível), execute o seguinte comando utilizando o ambiente python configurado com o ai_edge_litert:

# Script rápido de inspeção (inspect_litertlm.py)
import sys
from ai_edge_litert.internal import litertlm_core
from ai_edge_litert.internal import litertlm_header_schema_py_generated as schema

def inspect_bundle(caminho):
    dados = open(caminho, "rb").read()
    assert dados[:8] == b"LITERTLM", "Arquivo não é um bundle LiteRTLM válido!"
    
    header_end = int.from_bytes(dados[24:32], "little")
    meta = schema.LiteRTLMMetaData.GetRootAs(
        dados[litertlm_core.HEADER_BEGIN_BYTE_OFFSET:header_end], 0
    )
    sec = meta.SectionMetadata()
    
    print(f"📊 Bundle: {caminho} | Total de Seções: {sec.ObjectsLength()}")
    print("-" * 75)
    for i in range(sec.ObjectsLength()):
        obj = sec.Objects(i)
        tamanho = (obj.EndOffset() - obj.BeginOffset()) / (1024 * 1024)
        dtype = litertlm_core.any_section_data_type_to_string(obj.DataType())
        
        # Procura por model_type no vetor de itens
        model_type = "N/A"
        for j in range(obj.ItemsLength()):
            kv = obj.Items(j)
            if kv.Key() and kv.Key().decode() == "model_type":
                val = kv.Value()
                if val:
                    creator = schema.VdataCreator(kv.ValueType(), val)
                    val_bytes = getattr(creator, "value", None)
                    if isinstance(val_bytes, bytes):
                        model_type = val_bytes.decode()
        
        print(f"[{i:02d}] Tipo: {dtype:<25} | Model-Type: {model_type:<25} | Tam: {tamanho:>7.2f} MB")

if __name__ == "__main__":
    inspect_bundle(sys.argv[1])

---

<a id="pipeline-gemma-code"></a>

📝 7. Código-Fonte Completo de Pipeline_gemma.py

O script de compilação Pipeline_gemma.py é apresentado abaixo na íntegra. Ele automatiza o download do adaptador LoRA da Hugging Face no repositório ThiagoAllan/gemma-4-E2B-it-AgroRAG, realiza a fusão (merge) com o modelo base da Google, e exporta o bundle final utilizando a CLI litert-torch.

import os
import sys
import gc
import subprocess
import torch
from dotenv import load_dotenv
from huggingface_hub import login
from transformers import AutoModelForCausalLM, AutoTokenizer
from peft import PeftModel
import peft.tuners.lora.model as lora_tracker

# CONFIGURAÇÕES GLOBAIS
MODELO_BASE      = "google/gemma-4-E2B-it"
REPOSITORIO_LORA = "ThiagoAllan/gemma-4-E2B-it-AgroRAG"
SUBPASTA_LORA    = "lora"

CAMINHO_LOCAL    = "./gemma_4_fusion_local"
CAMINHO_LITERT   = "./gemma_4_RAG_mobile"

# Configura o PyTorch para usar os 36 cores da VM de forma eficiente
torch.set_num_threads(36)

# ==========================================
# FASE 0: AUTENTICAÇÃO
# ==========================================
print("\n" + "="*50)
print("FASE 0: AUTENTICAÇÃO")
print("="*50)

load_dotenv()
hf_token = os.getenv("HF_TOKEN") 

login(token=hf_token)
print("✅ Autenticado no HuggingFace!")
print("⚠️ Executando em modo CPU (Fusão de pesos + Conversão).")

# ==========================================
# FASE 1: CARREGAMENTO E FUSÃO (MERGE)
# ==========================================
print("\n" + "="*50)
print("FASE 1: CARREGANDO BASE + ADAPTADORES AGRO RAG")
print("="*50)

print(f"Carregando o Tokenizer oficial do modelo base...")
tokenizer = AutoTokenizer.from_pretrained(MODELO_BASE)
print("✅ Tokenizer carregado.")

print(f"Carregando modelo base da Google na RAM...")
modelo_base = AutoModelForCausalLM.from_pretrained(
    MODELO_BASE,
    device_map="cpu",
)
print("✅ Modelo base carregado.")

# --- PATCH DE VALIDAÇÃO PEFT PARA O MODELO BASE GEMMA 4 ---
_original_check_target_module_exists = getattr(lora_tracker, "_check_target_module_exists", None)
def _patched_check_target_module_exists(lora_config, key):
    if any(mod in key for mod in ["q_proj", "k_proj", "v_proj", "o_proj"]):
        return True
    if _original_check_target_module_exists:
        return _original_check_target_module_exists(lora_config, key)
    return True
lora_tracker._check_target_module_exists = _patched_check_target_module_exists
# -----------------------------------------------------------

print(f"Acoplado adaptadores LoRA do repositório '{REPOSITORIO_LORA}' na pasta '{SUBPASTA_LORA}'...")
modelo_peft = PeftModel.from_pretrained(
    modelo_base, 
    REPOSITORIO_LORA, 
    subfolder=SUBPASTA_LORA,
    device_map="cpu"
)
print("✅ Adaptadores AgroRAG acoplados com sucesso!")

print("Executando fusão dos pesos na RAM (Merge)...")
modelo_fundido = modelo_peft.merge_and_unload()

print(f"Salvando modelo completo fundido em '{CAMINHO_LOCAL}'...")
modelo_fundido.save_pretrained(CAMINHO_LOCAL, safe_serialization=True)
tokenizer.save_pretrained(CAMINHO_LOCAL)
print("✅ Modelo estruturado localmente e pronto para o conversor.")

# Limpeza agressiva de memória RAM
del modelo_base
del modelo_peft
del modelo_fundido
gc.collect()

# ==========================================
# FASE 2: CONVERSÃO LITERT VIA SUBPROCESS
# ==========================================
print("\n" + "="*50)
print("FASE 2: CONVERSÃO PARA MOBILE (LITERT-LM)")
print("="*50)

os.makedirs(CAMINHO_LITERT, exist_ok=True)
os.environ["PROTOCOL_BUFFERS_PYTHON_IMPLEMENTATION"] = "python"

# Tenta localizar 'litert-torch' no mesmo diretório do executável Python (ex: .venv/bin)
diretorio_bin = os.path.dirname(sys.executable)
executavel_litert = os.path.join(diretorio_bin, "litert-torch")
if not os.path.exists(executavel_litert):
    executavel_litert = "litert-torch"

comando_litert = [
    executavel_litert, "export_hf",
    f"--model={CAMINHO_LOCAL}",
    f"--output_dir={CAMINHO_LITERT}",
    "--task=text_generation",
    "--prefill_lengths=[32]",
    "--quantization_recipe=dynamic_wi4_afp32",
    "--externalize_embedder",
    "--jinja_chat_template_override=litert-community/gemma-4-E2B-it-litert-lm"
]

print(f"⏳ Executando conversão LiteRT: {' '.join(comando_litert)}")
resultado = subprocess.run(comando_litert, capture_output=False, text=True)

if resultado.returncode == 0:
    print(f"✅ Conversão concluída! Bundle em: {CAMINHO_LITERT}")
else:
    print(f"❌ Erro na conversão do LiteRT. Código de saída: {resultado.returncode}")

# ==========================================
# FASE 3: VERIFICAÇÃO DO BUNDLE FINAL
# ==========================================
print("\n" + "="*50)
print("FASE 3: VERIFICAÇÃO DO BUNDLE FINAL")
print("="*50)

if os.path.exists(CAMINHO_LITERT) and os.listdir(CAMINHO_LITERT):
    print(f"\n📁 Conteúdo de '{CAMINHO_LITERT}':\n")
    tamanho_total = 0
    for raiz, dirs, arquivos in os.walk(CAMINHO_LITERT):
        for arq in sorted(arquivos):
            caminho_completo = os.path.join(raiz, arq)
            tamanho = os.path.getsize(caminho_completo)
            tamanho_total += tamanho
            relativo = os.path.relpath(caminho_completo, CAMINHO_LITERT)
            print(f"  {relativo:<50}  {tamanho / 1e6:>8.1f} MB")

    print(f"\n{'-'*60}")
    print(f"  TOTAL{' '*45}{tamanho_total / 1e9:>8.2f} GB")
    print(f"\n🎉 Seu bundle AgroRAG está pronto para uso offline no App mobile!")
else:
    print("⚠️ A pasta do LiteRT está vazia. Verifique os logs de erro da Fase 2.")