Final tested Colab notebook with all fixes (correct weight naming, externalize_embedder, etc.)

gemma4_to_litertlm.ipynb

@@ -4,7 +4,8 @@
   "metadata": {
     "colab": {
       "provenance": [],
-      "gpuType": "T4"
+      "gpuType": "T4",
+      "machine_shape": "hm"
     },
     "kernelspec": {
       "name": "python3",
@@ -22,23 +23,11 @@
       "source": [
         "# 🚀 Convertir Gemma 4 E2B Uncensored-MAX a LiteRT-LM\n",
         "\n",
-        "Este notebook convierte [prithivMLmods/gemma-4-E2B-it-Uncensored-MAX](https://huggingface.co/prithivMLmods/gemma-4-E2B-it-Uncensored-MAX) al formato `.litertlm` para usarlo con **Google AI Edge Gallery** en Android.\n",
+        "Convierte el modelo a formato `.litertlm` para **Google AI Edge Gallery** en Android.\n",
         "\n",
-        "**Requisitos:**\n",
-        "- Colab con GPU (T4) — el runtime normal funciona, pero si te da OOM usa \"High-RAM\" (Entorno de ejecución → Cambiar tipo de entorno de ejecución → RAM Alta)\n",
-        "- Token de HuggingFace con permisos de escritura\n",
+        "**⚠️ IMPORTANTE:** Usa runtime con **GPU + RAM Alta**: Entorno de ejecución → Cambiar tipo → T4 + RAM Alta (hm)\n",
         "\n",
-        "**Tiempo estimado:** ~20-40 minutos"
-      ]
-    },
-    {
-      "cell_type": "markdown",
-      "metadata": {},
-      "source": [
-        "## 1️⃣ Configura tu token de HuggingFace\n",
-        "\n",
-        "Necesitas un token con permisos de escritura para subir el modelo. \n",
-        "Consíguelo en: https://huggingface.co/settings/tokens"
+        "**Tiempo estimado:** ~30-45 minutos"
       ]
     },
     {
@@ -47,24 +36,11 @@
       "metadata": {},
       "outputs": [],
       "source": [
-        "# ⬇️ PON TU TOKEN AQUÍ ⬇️\n",
-        "HF_TOKEN = \"\"  # Pega tu token de HuggingFace aquí (hf_...)\n",
-        "\n",
-        "# El repo donde se subirá el modelo convertido\n",
-        "# Cámbialo por tu usuario si quieres\n",
-        "OUTPUT_REPO = \"RedSparkie/gemma-4-E2B-it-Uncensored-MAX-litert-lm\"\n",
-        "\n",
-        "# Modelo fuente (los pesos originales en safetensors)\n",
-        "SOURCE_MODEL = \"prithivMLmods/gemma-4-E2B-it-Uncensored-MAX\"\n",
-        "\n",
-        "assert HF_TOKEN, \"❌ ¡Pon tu token de HuggingFace arriba!\""
-      ]
-    },
-    {
-      "cell_type": "markdown",
-      "metadata": {},
-      "source": [
-        "## 2️⃣ Instalar dependencias"
+        "#@title 1️⃣ Configuración\n",
+        "HF_TOKEN = \"\"  #@param {type:\"string\"}\n",
+        "OUTPUT_REPO = \"RedSparkie/gemma-4-E2B-it-Uncensored-MAX-litert-lm\"  #@param {type:\"string\"}\n",
+        "SOURCE_MODEL = \"prithivMLmods/gemma-4-E2B-it-Uncensored-MAX\"  #@param {type:\"string\"}\n",
+        "assert HF_TOKEN, \"❌ ¡Pon tu token de HuggingFace!\""
       ]
     },
     {
@@ -73,19 +49,9 @@
       "metadata": {},
       "outputs": [],
       "source": [
+        "#@title 2️⃣ Instalar dependencias\n",
         "!pip install -q litert-torch litert-lm transformers huggingface_hub sentencepiece protobuf safetensors psutil\n",
-        "print(\"✅ Dependencias instaladas\")"
-      ]
-    },
-    {
-      "cell_type": "markdown",
-      "metadata": {},
-      "source": [
-        "## 3️⃣ Extraer solo los pesos del decoder de texto\n",
-        "\n",
-        "El modelo original es multimodal (texto + visión + audio = 9.6 GB). \n",
-        "Nosotros solo necesitamos el decoder de texto (~4.8 GB en bf16). \n",
-        "Esto ahorra mucha RAM."
+        "print(\"✅ Instalado\")"
       ]
     },
     {
@@ -94,139 +60,106 @@
       "metadata": {},
       "outputs": [],
       "source": [
-        "import os, json, gc, shutil, sys, time\n",
+        "#@title 3️⃣ Preparar modelo (extraer solo texto, sin visión/audio)\n",
+        "import os, sys, json, gc, shutil, time\n",
         "from huggingface_hub import hf_hub_download\n",
         "from safetensors import safe_open\n",
         "from safetensors.torch import save_file\n",
-        "import torch\n",
-        "import transformers\n",
-        "import psutil\n",
+        "import transformers, psutil\n",
         "\n",
-        "def memlog(label=\"\"):\n",
-        "    m = psutil.virtual_memory()\n",
-        "    print(f\"  [{label}] RAM: {m.available/(1024**3):.1f}/{m.total/(1024**3):.1f} GB disponibles\")\n",
+        "def memlog(l=\"\"):\n",
+        "    m=psutil.virtual_memory()\n",
+        "    print(f\"  [{l}] RAM: {m.available/(1024**3):.1f}/{m.total/(1024**3):.1f} GB\")\n",
         "\n",
-        "TEXT_MODEL_DIR = \"/content/text_model\"\n",
+        "MODEL_DIR = \"/content/model\"\n",
         "OUTPUT_DIR = \"/content/output\"\n",
-        "os.makedirs(TEXT_MODEL_DIR, exist_ok=True)\n",
+        "os.makedirs(MODEL_DIR, exist_ok=True)\n",
         "os.makedirs(OUTPUT_DIR, exist_ok=True)\n",
-        "\n",
         "start_time = time.time()\n",
-        "memlog(\"INICIO\")\n",
+        "memlog(\"inicio\")\n",
         "\n",
-        "# Descargar index de shards\n",
-        "print(\"📥 Descargando índice de pesos...\")\n",
-        "idx_file = hf_hub_download(SOURCE_MODEL, \"model.safetensors.index.json\", token=HF_TOKEN)\n",
-        "with open(idx_file) as f:\n",
+        "# Descargar índice\n",
+        "print(\"📥 Descargando índice...\")\n",
+        "idx_path = hf_hub_download(SOURCE_MODEL, \"model.safetensors.index.json\", token=HF_TOKEN)\n",
+        "with open(idx_path) as f:\n",
         "    index = json.load(f)\n",
         "\n",
-        "# Identificar shards que contienen pesos del language model\n",
-        "shard_keys = {}\n",
+        "# Agrupar pesos del language model por shard\n",
+        "shard_lm = {}\n",
         "for key, shard in index[\"weight_map\"].items():\n",
         "    if key.startswith(\"model.language_model.\"):\n",
-        "        if shard not in shard_keys:\n",
-        "            shard_keys[shard] = []\n",
-        "        shard_keys[shard].append(key)\n",
-        "\n",
-        "print(f\"  Encontrados {sum(len(v) for v in shard_keys.values())} tensores de texto en {len(shard_keys)} shards\")\n",
+        "        shard_lm.setdefault(shard, []).append(key)\n",
         "\n",
-        "# Procesar shard por shard: extraer solo pesos LM, guardar y liberar\n",
-        "new_weight_map = {}\n",
-        "shard_idx = 0\n",
+        "total_shards = len(shard_lm)\n",
+        "print(f\"  {sum(len(v) for v in shard_lm.values())} tensores en {total_shards} shards\")\n",
         "\n",
-        "for shard_name in sorted(shard_keys.keys()):\n",
-        "    keys_in_shard = shard_keys[shard_name]\n",
-        "    print(f\"\\n📦 Procesando {shard_name} ({len(keys_in_shard)} tensores)...\")\n",
+        "# Extraer shard por shard (MANTENER el prefijo model.language_model.)\n",
+        "weight_map = {}\n",
+        "for i, sn in enumerate(sorted(shard_lm)):\n",
+        "    keys = shard_lm[sn]\n",
+        "    out_name = f\"model-{i+1:05d}-of-{total_shards:05d}.safetensors\"\n",
+        "    out_path = os.path.join(MODEL_DIR, out_name)\n",
         "    \n",
-        "    # Descargar shard\n",
-        "    shard_path = hf_hub_download(SOURCE_MODEL, shard_name, token=HF_TOKEN)\n",
+        "    if os.path.exists(out_path) and os.path.getsize(out_path) > 100:\n",
+        "        print(f\"  {out_name} ya existe, skip\")\n",
+        "        with safe_open(out_path, framework=\"pt\") as f:\n",
+        "            for k in f.keys(): weight_map[k] = out_name\n",
+        "        continue\n",
         "    \n",
-        "    # Extraer solo tensores del language_model, quitar prefijo\n",
-        "    lm_weights = {}\n",
-        "    with safe_open(shard_path, framework=\"pt\") as f:\n",
-        "        for key in keys_in_shard:\n",
-        "            new_key = key[len(\"model.language_model.\"):]\n",
-        "            lm_weights[new_key] = f.get_tensor(key)\n",
+        "    print(f\"  📦 {sn} → {out_name} ({len(keys)} tensores)\")\n",
+        "    shard_path = hf_hub_download(SOURCE_MODEL, sn, token=HF_TOKEN)\n",
         "    \n",
-        "    # Guardar como nuevo shard\n",
-        "    shard_idx += 1\n",
-        "    out_name = f\"model-{shard_idx:05d}-of-TEMP.safetensors\"\n",
-        "    out_path = os.path.join(TEXT_MODEL_DIR, out_name)\n",
-        "    save_file(lm_weights, out_path)\n",
+        "    # Extraer tensores MANTENIENDO el prefijo original\n",
+        "    tensors = {}\n",
+        "    with safe_open(shard_path, framework=\"pt\") as f:\n",
+        "        for key in keys:\n",
+        "            tensors[key] = f.get_tensor(key)\n",
         "    \n",
-        "    for k in lm_weights:\n",
-        "        new_weight_map[k] = out_name\n",
+        "    save_file(tensors, out_path)\n",
+        "    for k in tensors: weight_map[k] = out_name\n",
         "    \n",
         "    size_mb = os.path.getsize(out_path) / (1024**2)\n",
-        "    print(f\"  💾 Guardado {out_name}: {size_mb:.0f} MB\")\n",
-        "    \n",
-        "    del lm_weights\n",
-        "    gc.collect()\n",
-        "    memlog(f\"shard {shard_idx}\")\n",
-        "\n",
-        "# Renombrar shards con total correcto\n",
-        "total_shards = shard_idx\n",
-        "final_weight_map = {}\n",
-        "for i in range(1, total_shards + 1):\n",
-        "    old_name = f\"model-{i:05d}-of-TEMP.safetensors\"\n",
-        "    new_name = f\"model-{i:05d}-of-{total_shards:05d}.safetensors\"\n",
-        "    os.rename(os.path.join(TEXT_MODEL_DIR, old_name), os.path.join(TEXT_MODEL_DIR, new_name))\n",
-        "    for key, shard in new_weight_map.items():\n",
-        "        if shard == old_name:\n",
-        "            final_weight_map[key] = new_name\n",
+        "    print(f\"    💾 {size_mb:.0f} MB\")\n",
+        "    del tensors; gc.collect()\n",
+        "    memlog(f\"shard {i+1}\")\n",
         "\n",
         "# Escribir índice\n",
-        "with open(os.path.join(TEXT_MODEL_DIR, \"model.safetensors.index.json\"), \"w\") as f:\n",
-        "    json.dump({\"metadata\": {}, \"weight_map\": final_weight_map}, f)\n",
+        "with open(os.path.join(MODEL_DIR, \"model.safetensors.index.json\"), \"w\") as f:\n",
+        "    json.dump({\"metadata\": {}, \"weight_map\": weight_map}, f)\n",
         "\n",
-        "# Config: usar Gemma4TextConfig como standalone\n",
+        "# Config: Gemma4 con vision=None, audio=None\n",
+        "print(\"\\n📝 Creando config...\")\n",
         "config = transformers.AutoConfig.from_pretrained(SOURCE_MODEL, token=HF_TOKEN)\n",
-        "tc = config.text_config.to_dict()\n",
-        "tc[\"architectures\"] = [\"Gemma4ForCausalLM\"]\n",
-        "tc[\"model_type\"] = \"gemma4_text\"\n",
-        "tc[\"eos_token_id\"] = config.eos_token_id if hasattr(config, \"eos_token_id\") else [1, 106]\n",
-        "tc[\"tie_word_embeddings\"] = config.tie_word_embeddings\n",
-        "with open(os.path.join(TEXT_MODEL_DIR, \"config.json\"), \"w\") as f:\n",
-        "    json.dump(tc, f, indent=2)\n",
-        "\n",
-        "# Copiar tokenizer y templates\n",
+        "cd = config.to_dict()\n",
+        "cd[\"vision_config\"] = None\n",
+        "cd[\"audio_config\"] = None\n",
+        "for k in [\"vision_soft_tokens_per_image\", \"image_token_id\", \"boi_token_id\",\n",
+        "          \"eoi_token_id\", \"audio_token_id\", \"boa_token_id\", \"eoa_token_id\",\n",
+        "          \"eoa_token_index\", \"video_token_id\"]:\n",
+        "    cd.pop(k, None)\n",
+        "with open(os.path.join(MODEL_DIR, \"config.json\"), \"w\") as f:\n",
+        "    json.dump(cd, f, indent=2)\n",
+        "\n",
+        "# Tokenizer y archivos extra\n",
         "for fn in [\"tokenizer.json\", \"tokenizer_config.json\", \"chat_template.jinja\", \"generation_config.json\"]:\n",
         "    try:\n",
         "        src = hf_hub_download(SOURCE_MODEL, fn, token=HF_TOKEN)\n",
-        "        shutil.copy(src, os.path.join(TEXT_MODEL_DIR, fn))\n",
-        "        print(f\"  📄 {fn}\")\n",
-        "    except:\n",
-        "        pass\n",
-        "\n",
-        "del config\n",
+        "        shutil.copy(src, os.path.join(MODEL_DIR, fn))\n",
+        "        print(f\"  ✓ {fn}\")\n",
+        "    except: pass\n",
+        "\n",
+        "del config; gc.collect()\n",
+        "\n",
+        "# Limpiar caché HF\n",
+        "cache_dir = os.path.expanduser(\"~/.cache/huggingface/hub\")\n",
+        "if os.path.exists(cache_dir):\n",
+        "    for d in os.listdir(cache_dir):\n",
+        "        if d.startswith(\"models--\"):\n",
+        "            shutil.rmtree(os.path.join(cache_dir, d), ignore_errors=True)\n",
         "gc.collect()\n",
         "\n",
-        "# Resumen\n",
-        "print(f\"\\n✅ Modelo de texto extraído en {TEXT_MODEL_DIR}:\")\n",
-        "total_size = 0\n",
-        "for f in sorted(os.listdir(TEXT_MODEL_DIR)):\n",
-        "    fp = os.path.join(TEXT_MODEL_DIR, f)\n",
-        "    if os.path.isfile(fp):\n",
-        "        s = os.path.getsize(fp)\n",
-        "        total_size += s\n",
-        "        print(f\"  {f}: {s/(1024**2):.1f} MB\")\n",
-        "print(f\"  Total: {total_size/(1024**3):.2f} GB\")\n",
-        "print(f\"  Tiempo: {(time.time()-start_time)/60:.1f} min\")"
-      ]
-    },
-    {
-      "cell_type": "markdown",
-      "metadata": {},
-      "source": [
-        "## 4️⃣ Convertir a LiteRT-LM (.litertlm)\n",
-        "\n",
-        "Aquí es donde ocurre la magia. El pipeline de `litert-torch` hace:\n",
-        "1. Cargar el modelo en float32\n",
-        "2. Exportar a TFLite via torch.export\n",
-        "3. Cuantizar a INT8 (dynamic_wi8_afp32)\n",
-        "4. Empaquetar como `.litertlm`\n",
-        "\n",
-        "⚠️ **Si te da error de memoria**, ve a: Entorno de ejecución → Cambiar tipo → **RAM Alta**"
+        "print(f\"\\n✅ Modelo preparado\")\n",
+        "memlog(\"listo\")"
       ]
     },
     {
@@ -235,99 +168,33 @@
       "metadata": {},
       "outputs": [],
       "source": [
-        "import litert_torch.generative.export_hf.core.export_lib as elib\n",
-        "from litert_torch.generative.export_hf.model_ext import patches as mpatches\n",
-        "from litert_torch.generative.export_hf.core import exportable_module_config\n",
-        "from litert_torch.generative.export_hf.core import utils\n",
-        "from litert_torch import progress\n",
-        "import huggingface_hub as hfhub\n",
-        "\n",
-        "ExportTask = exportable_module_config.ExportTask\n",
-        "\n",
-        "# Monkey-patch load_model para cargar Gemma4ForCausalLM desde nuestro dir\n",
-        "@progress.task('Load source model')\n",
-        "def patched_load_model(model_path, trust_remote_code=False, auto_model_override=None, task=ExportTask.TEXT_GENERATION):\n",
-        "    print(\"  🔧 Cargando Gemma4ForCausalLM (solo texto)...\")\n",
-        "    \n",
-        "    config = transformers.AutoConfig.from_pretrained(model_path, trust_remote_code=trust_remote_code)\n",
-        "    config.model_type = \"gemma4\"  # Para que el pipeline reconozca la arquitectura\n",
-        "    config._attn_implementation = 'lrt_transposed_attention'\n",
-        "    \n",
-        "    from transformers import Gemma4ForCausalLM\n",
-        "    with mpatches.get_patch_context(\"gemma4\"):\n",
-        "        model = Gemma4ForCausalLM.from_pretrained(\n",
-        "            model_path, config=config, torch_dtype=torch.float32,\n",
-        "            trust_remote_code=trust_remote_code, low_cpu_mem_usage=True,\n",
-        "            attn_implementation='eager',\n",
-        "        )\n",
-        "    \n",
-        "    memlog(\"modelo cargado\")\n",
-        "    \n",
-        "    model.generation_config.cache_implementation = 'static'\n",
-        "    model.generation_config.do_sample = False\n",
-        "    \n",
-        "    # El pipeline espera un config con text_config\n",
-        "    class FullConfig:\n",
-        "        def __init__(self, tc):\n",
-        "            self.model_type = \"gemma4\"\n",
-        "            self.text_config = tc\n",
-        "            self.eos_token_id = getattr(tc, 'eos_token_id', [1, 106])\n",
-        "            self.tie_word_embeddings = getattr(tc, 'tie_word_embeddings', True)\n",
-        "    \n",
-        "    full_config = FullConfig(config)\n",
-        "    model.config = full_config\n",
-        "    \n",
-        "    tokenizer = transformers.AutoTokenizer.from_pretrained(model_path)\n",
-        "    \n",
-        "    # Cargar chat template\n",
-        "    if not getattr(tokenizer, 'chat_template', None):\n",
-        "        jinja_path = os.path.join(model_path, 'chat_template.jinja')\n",
-        "        if os.path.exists(jinja_path):\n",
-        "            with open(jinja_path) as f:\n",
-        "                tokenizer.chat_template = f.read()\n",
-        "    \n",
-        "    return elib.SourceModelArtifacts(\n",
-        "        model=model, model_config=full_config,\n",
-        "        text_model_config=config, tokenizer=tokenizer, image_processor=None,\n",
-        "    )\n",
-        "\n",
-        "# Aplicar el patch\n",
-        "elib.load_model = patched_load_model\n",
-        "\n",
-        "# Ejecutar la conversión\n",
+        "#@title 4️⃣ Convertir a .litertlm\n",
+        "import torch\n",
         "from litert_torch.generative.export_hf import export as export_lib\n",
         "\n",
-        "print(\"🚀 Iniciando conversión a LiteRT-LM...\")\n",
-        "print(f\"  Fuente: {TEXT_MODEL_DIR}\")\n",
-        "print(f\"  Destino: {OUTPUT_DIR}\")\n",
-        "print(f\"  Cuantización: dynamic_wi8_afp32 (INT8)\")\n",
-        "print(f\"  Cache: 4096 tokens\")\n",
-        "print()\n",
+        "print(\"🚀 Convirtiendo a LiteRT-LM...\")\n",
+        "print(\"  Esto tarda 15-30 min. Paciencia.\")\n",
+        "memlog(\"pre-export\")\n",
         "\n",
         "conversion_start = time.time()\n",
         "\n",
         "export_lib.export(\n",
-        "    model=TEXT_MODEL_DIR,\n",
+        "    model=MODEL_DIR,\n",
         "    output_dir=OUTPUT_DIR,\n",
         "    task=\"text_generation\",\n",
         "    bundle_litert_lm=True,\n",
-        "    quantization_recipe=\"dynamic_wi8_afp32\",\n",
+        "    quantization_recipe=\"dynamic_wi8_afp32\",  # INT8 (como el oficial)\n",
         "    cache_length=4096,\n",
         "    prefill_lengths=[256],\n",
         "    use_jinja_template=True,\n",
         "    keep_temporary_files=True,\n",
         "    trust_remote_code=False,\n",
         "    experimental_lightweight_conversion=True,\n",
+        "    externalize_embedder=True,  # Requerido para Gemma4\n",
         ")\n",
         "\n",
-        "print(f\"\\n✅ Conversión completada en {(time.time()-conversion_start)/60:.1f} minutos\")"
-      ]
-    },
-    {
-      "cell_type": "markdown",
-      "metadata": {},
-      "source": [
-        "## 5️⃣ Verificar y subir a HuggingFace"
+        "print(f\"\\n✅ Conversión en {(time.time()-conversion_start)/60:.1f} min\")\n",
+        "memlog(\"post-export\")"
       ]
     },
     {
@@ -336,57 +203,44 @@
       "metadata": {},
       "outputs": [],
       "source": [
-        "litertlm_path = os.path.join(OUTPUT_DIR, \"model.litertlm\")\n",
-        "\n",
-        "if not os.path.exists(litertlm_path):\n",
-        "    print(\"❌ model.litertlm no encontrado. Archivos generados:\")\n",
-        "    for root, dirs, files in os.walk(OUTPUT_DIR):\n",
-        "        for f in files:\n",
-        "            fp = os.path.join(root, f)\n",
-        "            print(f\"  {os.path.relpath(fp, OUTPUT_DIR)}: {os.path.getsize(fp)/(1024**2):.1f} MB\")\n",
+        "#@title 5️⃣ Verificar y subir\n",
+        "litertlm = os.path.join(OUTPUT_DIR, \"model.litertlm\")\n",
+        "\n",
+        "if not os.path.exists(litertlm):\n",
+        "    print(\"❌ model.litertlm no encontrado. Archivos:\")\n",
+        "    for r,d,fs in os.walk(OUTPUT_DIR):\n",
+        "        for f in fs:\n",
+        "            fp = os.path.join(r,f)\n",
+        "            print(f\"  {os.path.relpath(fp,OUTPUT_DIR)}: {os.path.getsize(fp)/(1024**2):.1f} MB\")\n",
         "else:\n",
-        "    size_bytes = os.path.getsize(litertlm_path)\n",
-        "    size_gb = size_bytes / (1024**3)\n",
-        "    print(f\"📊 model.litertlm: {size_gb:.2f} GB ({size_bytes:,} bytes)\")\n",
+        "    size_gb = os.path.getsize(litertlm) / (1024**3)\n",
+        "    print(f\"📊 model.litertlm: {size_gb:.2f} GB\")\n",
         "    if size_gb <= 2.0:\n",
-        "        print(f\"✅ ¡Cabe en 2 GB!\")\n",
+        "        print(\"✅ ¡Cabe en 2 GB!\")\n",
         "    else:\n",
-        "        print(f\"⚠️ Pesa más de 2 GB ({size_gb:.2f} GB)\")\n",
+        "        print(f\"⚠️ {size_gb:.2f} GB — Si necesitas menos, cambia a dynamic_wi4_afp32 en celda 4\")\n",
         "    \n",
         "    print(f\"\\n📤 Subiendo a {OUTPUT_REPO}...\")\n",
         "    from huggingface_hub import HfApi\n",
         "    api = HfApi(token=HF_TOKEN)\n",
+        "    try: api.create_repo(OUTPUT_REPO, exist_ok=True)\n",
+        "    except: pass\n",
         "    \n",
-        "    # Crear repo si no existe\n",
-        "    try:\n",
-        "        api.create_repo(OUTPUT_REPO, exist_ok=True)\n",
-        "    except:\n",
-        "        pass\n",
-        "    \n",
-        "    # Subir modelo\n",
         "    api.upload_file(\n",
-        "        path_or_fileobj=litertlm_path,\n",
+        "        path_or_fileobj=litertlm,\n",
         "        path_in_repo=\"gemma-4-E2B-it-Uncensored-MAX.litertlm\",\n",
         "        repo_id=OUTPUT_REPO,\n",
         "        commit_message=\"Add LiteRT-LM model\",\n",
         "    )\n",
         "    \n",
-        "    # Subir README\n",
-        "    readme = f\"\"\"---\\nlicense: apache-2.0\\nbase_model:\\n- prithivMLmods/gemma-4-E2B-it-Uncensored-MAX\\ntags:\\n  - litert-lm\\n  - uncensored\\n  - abliterated\\n  - edge-gallery\\n  - on-device\\nlanguage:\\n- en\\n---\\n\\n# gemma-4-E2B-it-Uncensored-MAX (LiteRT-LM)\\n\\nLiteRT-LM conversion of [prithivMLmods/gemma-4-E2B-it-Uncensored-MAX](https://huggingface.co/prithivMLmods/gemma-4-E2B-it-Uncensored-MAX) for **Google AI Edge Gallery** on Android.\\n\\n| | |\\n|---|---|\\n| **Base model** | [prithivMLmods/gemma-4-E2B-it-Uncensored-MAX](https://huggingface.co/prithivMLmods/gemma-4-E2B-it-Uncensored-MAX) |\\n| **Format** | LiteRT-LM (`.litertlm`) |\\n| **Quantization** | INT8 (`dynamic_wi8_afp32`) |\\n| **Task** | Text generation |\\n| **Context** | 4096 tokens |\\n| **Size** | {size_gb:.2f} GB |\\n\\n## Usage\\n\\n### Edge Gallery (Android)\\n1. Install [Google AI Edge Gallery](https://play.google.com/store/apps/details?id=com.google.ai.edge.gallery)\\n2. Add model via HuggingFace URL\\n3. Chat!\\n\\n### CLI\\n```bash\\npip install litert-lm\\nlitert-lm import --from-huggingface-repo {OUTPUT_REPO} gemma-4-E2B-it-Uncensored-MAX.litertlm uncensored-max\\nlitert-lm run uncensored-max\\n```\\n\\n⚠️ Abliterated/uncensored model. Use responsibly.\\n\"\"\"    \n",
-        "    api.upload_file(\n",
-        "        path_or_fileobj=readme.encode(),\n",
-        "        path_in_repo=\"README.md\",\n",
-        "        repo_id=OUTPUT_REPO,\n",
-        "        commit_message=\"Add README\",\n",
-        "    )\n",
+        "    readme = f\"\"\"---\\nlicense: apache-2.0\\nbase_model:\\n- {SOURCE_MODEL}\\ntags:\\n  - litert-lm\\n  - uncensored\\n  - edge-gallery\\nlanguage:\\n- en\\n---\\n\\n# gemma-4-E2B-it-Uncensored-MAX (LiteRT-LM)\\n\\nLiteRT-LM conversion for **Google AI Edge Gallery**.\\n\\n| | |\\n|---|---|\\n| **Base** | [{SOURCE_MODEL}](https://huggingface.co/{SOURCE_MODEL}) |\\n| **Format** | `.litertlm` |\\n| **Quant** | INT8 |\\n| **Context** | 4096 |\\n| **Size** | {size_gb:.2f} GB |\\n\\n## Usage\\n1. Install [Edge Gallery](https://play.google.com/store/apps/details?id=com.google.ai.edge.gallery)\\n2. Add model via HF URL\\n3. Chat!\\n\\n⚠️ Uncensored. Use responsibly.\\n\"\"\"\n",
+        "    api.upload_file(path_or_fileobj=readme.encode(), path_in_repo=\"README.md\",\n",
+        "                    repo_id=OUTPUT_REPO, commit_message=\"README\")\n",
         "    \n",
-        "    total_time = (time.time() - start_time) / 60\n",
-        "    print(f\"\\n{'='*50}\")\n",
-        "    print(f\"🎉 ¡LISTO!\")\n",
-        "    print(f\"📱 Modelo: https://huggingface.co/{OUTPUT_REPO}\")\n",
-        "    print(f\"📊 Tamaño: {size_gb:.2f} GB\")\n",
-        "    print(f\"⏱️ Tiempo total: {total_time:.0f} minutos\")\n",
-        "    print(f\"{'='*50}\")"
+        "    print(f\"\\n🎉 ¡LISTO!\")\n",
+        "    print(f\"📱 https://huggingface.co/{OUTPUT_REPO}\")\n",
+        "    print(f\"📊 {size_gb:.2f} GB\")\n",
+        "    print(f\"⏱️ {(time.time()-start_time)/60:.0f} min total\")"
       ]
     },
     {
@@ -395,17 +249,11 @@
       "source": [
         "## 🔧 Troubleshooting\n",
         "\n",
-        "**Error de memoria (OOM):**\n",
-        "- Ve a **Entorno de ejecución → Cambiar tipo de entorno** → Activa **RAM Alta**\n",
-        "- Si aún falla, reinicia el runtime y ejecuta todas las celdas de nuevo\n",
+        "**OOM:** Usa runtime con **RAM Alta** (hm)\n",
         "\n",
-        "**Error de `attn_implementation`:**\n",
-        "- Esto es normal si hay incompatibilidad entre versiones de `transformers` y `litert-torch`\n",
-        "- Intenta: `!pip install transformers==5.7.0`\n",
+        "**>2 GB:** Cambia `dynamic_wi8_afp32` → `dynamic_wi4_afp32` en celda 4\n",
         "\n",
-        "**El modelo pesa >2 GB:**\n",
-        "- Cambia `quantization_recipe` a `\"dynamic_wi4_afp32\"` (INT4) en la celda 4\n",
-        "- Esto reducirá el tamaño a la mitad pero con algo menos de calidad"
+        "**Error `External embedder required`:** Ya está solucionado con `externalize_embedder=True`"
       ]
     }
   ]