added gemma local

jupyter_notebooks/Section_2-3-4_Figure_8_deepfake_adapters.ipynb

@@ -2270,6 +2270,14 @@
    "metadata": {},
    "outputs": [],
    "source": []
+  },
+  {
+   "cell_type": "code",
+   "id": "4gb2vta3w5g",
+   "source": "import pandas as pd\nimport json\nimport time\nimport re\nfrom pathlib import Path\nfrom tqdm import tqdm\nimport torch\nfrom transformers import AutoModelForCausalLM, AutoTokenizer\n\ncurrent_dir = Path.cwd()\ninput_file = current_dir.parent / \"data/CSV/model_adapter/real_person_adapter_step_02_NER.csv\"\n\n# === CONFIGURATION ===\nTEST_MODE = True\nTEST_SIZE = 10\nMAX_ROWS = 20000\nSAVE_INTERVAL = 10\n\n# Model settings\nMODEL_NAME = \"google/gemma-3-27b-it\"\nCACHE_DIR = current_dir.parent / \"data/models\"\nCACHE_DIR.mkdir(parents=True, exist_ok=True)\n\nPROFESSION_CATEGORIES = [\n    \"actor\", \"adult performer\", \"singer/musician\", \"model\",\n    \"online personality\", \"public figure\", \"voice actor/ASMR\",\n    \"sports professional\", \"tv personality\"\n]\n\n# === LOAD MODEL ===\nprint(f\"Loading model: {MODEL_NAME}\")\nprint(f\"Cache directory: {CACHE_DIR}\")\nprint(f\"This may take a while on first run (~54GB download)...\\n\")\n\n# Check GPU availability\ndevice = \"cuda\" if torch.cuda.is_available() else \"cpu\"\nprint(f\"Device: {device}\")\n\nif device == \"cpu\":\n    print(\"⚠️  WARNING: No GPU detected! Inference will be VERY slow.\")\n    print(\"   Consider using a GPU or reducing model size.\")\n\n# Load tokenizer\nprint(\"Loading tokenizer...\")\ntokenizer = AutoTokenizer.from_pretrained(\n    MODEL_NAME,\n    cache_dir=str(CACHE_DIR),\n    trust_remote_code=True\n)\nprint(\"✅ Tokenizer loaded\")\n\n# Load model with optimizations\nprint(\"Loading model (this may take several minutes)...\")\nmodel = AutoModelForCausalLM.from_pretrained(\n    MODEL_NAME,\n    cache_dir=str(CACHE_DIR),\n    torch_dtype=torch.bfloat16,  # Use BF16 for efficiency\n    device_map=\"auto\",  # Automatically distribute across GPUs\n    trust_remote_code=True\n)\nmodel.eval()  # Set to evaluation mode\nprint(\"✅ Model loaded\")\n\n# Check VRAM usage\nif torch.cuda.is_available():\n    vram_gb = torch.cuda.max_memory_allocated() / 1024**3\n    print(f\"VRAM used: {vram_gb:.2f} GB\\n\")\n\n# === LOAD DATA ===\ndf = pd.read_csv(input_file)\nprint(f\"Loaded {len(df)} rows\")\n\nif TEST_MODE:\n    print(f\"Running in TEST MODE with {TEST_SIZE} samples\")\n    df = df.head(TEST_SIZE).copy()\nelif MAX_ROWS:\n    df = df.head(MAX_ROWS).copy()\n\n# === CREATE PROMPT ===\ndef create_prompt(name: str, tags: list, country: str) -> str:\n    tags_str = \", \".join(tags[:5]) if tags else \"none\"\n    country_str = country if country else \"unknown\"\n    \n    return f\"\"\"You are an expert at identifying public figures.\n\nGiven: Name: {name}, Tags: {tags_str}, Country: {country_str}\n\nProvide in JSON format:\n{{\n  \"full_name\": \"<full name>\",\n  \"gender\": \"<Male/Female/Other/Unknown>\",\n  \"professions\": [\"<profession 1>\", \"<profession 2>\", \"<profession 3>\"],\n  \"country\": \"<country>\"\n}}\n\nChoose professions ONLY from: {\", \".join(PROFESSION_CATEGORIES)}\n- Max 3 professions, ordered by relevance\n- \"online personality\" includes streamers, cosplayers, influencers\n- Return \"Unknown\" if uncertain\n\nReturn ONLY JSON, no extra text.\"\"\"\n\n# === QUERY GEMMA LOCAL ===\ndef query_gemma_local(prompt: str) -> dict:\n    \"\"\"Query Gemma-3-27B locally via transformers.\"\"\"\n    try:\n        # Format as chat message\n        messages = [\n            {\"role\": \"user\", \"content\": prompt}\n        ]\n        \n        # Tokenize\n        text = tokenizer.apply_chat_template(\n            messages,\n            tokenize=False,\n            add_generation_prompt=True\n        )\n        \n        inputs = tokenizer([text], return_tensors=\"pt\").to(device)\n        \n        # Generate\n        with torch.no_grad():\n            outputs = model.generate(\n                **inputs,\n                max_new_tokens=512,\n                temperature=0.1,\n                do_sample=True,\n                top_p=0.9,\n                pad_token_id=tokenizer.eos_token_id\n            )\n        \n        # Decode\n        generated_ids = outputs[0][inputs['input_ids'].shape[1]:]\n        response = tokenizer.decode(generated_ids, skip_special_tokens=True)\n        \n        # Parse JSON from response\n        json_match = re.search(r'\\{[^}]+\\}', response, re.DOTALL)\n        if json_match:\n            return json.loads(json_match.group())\n        \n        # Try markdown code block\n        code_block = re.search(r'```(?:json)?\\s*(\\{[^}]+\\})\\s*```', response, re.DOTALL)\n        if code_block:\n            return json.loads(code_block.group(1))\n        \n        return None\n        \n    except json.JSONDecodeError as e:\n        print(f\"JSON parse error: {e}\")\n        return None\n    except Exception as e:\n        print(f\"Generation error: {e}\")\n        return None\n\n# === PROCESS ===\noutput_file = current_dir.parent / f\"data/CSV/gemma_local_annotated_POI{'_test' if TEST_MODE else ''}.csv\"\nindex_file = current_dir.parent / \"misc/gemma_local_query_index.txt\"\n\ncurrent_index = 0\nif index_file.exists():\n    with open(index_file) as f:\n        current_index = int(f.read().strip())\n    print(f\"Resuming from index {current_index}\")\n\nfor col in ['full_name', 'gender', 'profession_llm', 'country']:\n    if col not in df.columns:\n        df[col] = 'Unknown'\n\nprint(f\"\\nAnnotating with GEMMA-3-27B LOCAL - rows {current_index} to {len(df)}...\")\nprint(f\"Model: {MODEL_NAME}\")\nprint(f\"This may take a while (~250-600 samples/hour on A100)\\n\")\n\ntry:\n    start_time = time.time()\n    \n    for i in tqdm(range(current_index, len(df)), desc=\"Gemma Local\"):\n        row = df.iloc[i]\n        \n        tags = [str(row[f'tag_{j}']).lower() for j in range(1, 8)\n                if f'tag_{j}' in row and pd.notna(row[f'tag_{j}'])]\n        \n        prompt = create_prompt(\n            row.get('real_name', row.get('name', '')),\n            tags,\n            row.get('likely_country', '')\n        )\n        \n        result = query_gemma_local(prompt)\n        \n        if result:\n            df.at[i, 'full_name'] = result.get('full_name', 'Unknown')\n            df.at[i, 'gender'] = result.get('gender', 'Unknown')\n            profs = result.get('professions', [])\n            df.at[i, 'profession_llm'] = \", \".join(profs[:3]) if isinstance(profs, list) else str(profs)\n            df.at[i, 'country'] = result.get('country', 'Unknown')\n        \n        current_index = i + 1\n        \n        if (i + 1) % SAVE_INTERVAL == 0 or (i + 1) == len(df):\n            df.to_csv(output_file, index=False)\n            with open(index_file, 'w') as f:\n                f.write(str(current_index))\n            \n            # Calculate stats\n            elapsed = time.time() - start_time\n            samples_done = current_index - (0 if not index_file.exists() else int(open(index_file).read().strip()))\n            samples_per_hour = (samples_done / elapsed) * 3600 if elapsed > 0 else 0\n            \n            print(f\"✅ Saved after {i+1} rows (~{samples_per_hour:.1f} samples/hour)\")\n    \n    elapsed_total = time.time() - start_time\n    print(f\"\\n✅ Done! Results: {output_file}\")\n    print(f\"Total time: {elapsed_total/60:.1f} minutes\")\n    print(f\"Average speed: {len(df)/(elapsed_total/3600):.1f} samples/hour\")\n    \n    print(f\"\\nTop professions:\\n{df['profession_llm'].value_counts().head(10)}\")\n    \n    # Show VRAM stats\n    if torch.cuda.is_available():\n        print(f\"\\nFinal VRAM usage: {torch.cuda.max_memory_allocated() / 1024**3:.2f} GB\")\n\nexcept KeyboardInterrupt:\n    print(f\"\\n⚠️ Interrupted by user\")\n    df.to_csv(output_file, index=False)\n    with open(index_file, 'w') as f:\n        f.write(str(current_index))\n    print(f\"Progress saved up to row {current_index}\")\n\nexcept Exception as e:\n    print(f\"⚠️ Error: {e}\")\n    df.to_csv(output_file, index=False)\n    with open(index_file, 'w') as f:\n        f.write(str(current_index))\n    print(f\"Progress saved up to row {current_index}\")",
+   "metadata": {},
+   "execution_count": null,
+   "outputs": []
   }
  ],
  "metadata": {
@@ -2293,4 +2301,4 @@
  },
  "nbformat": 4,
  "nbformat_minor": 5
-}
+}