added looging to bloom query

jupyter_notebooks/Section_2-3-4_Bloomz_query.ipynb

@@ -7,342 +7,408 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "import pandas as pd\n",
-    "import json\n",
-    "import time\n",
-    "import re\n",
-    "from pathlib import Path\n",
-    "from tqdm import tqdm\n",
-    "import torch\n",
-    "from transformers import AutoModelForCausalLM, AutoTokenizer\n",
-    "\n",
-    "# Import is used for pd.notna() and pd.isna() checks\n",
-    "\n",
-    "current_dir = Path.cwd()\n",
-    "input_file = current_dir.parent / \"data/CSV/model_adapter/real_person_adapter_step_02_NER.csv\"\n",
-    "\n",
-    "# === CONFIGURATION ===\n",
-    "TEST_MODE = True\n",
-    "TEST_SIZE = 10\n",
-    "MAX_ROWS = 20000\n",
-    "SAVE_INTERVAL = 10\n",
-    "\n",
-    "# Model settings - BLOOMZ (BigScience - European consortium)\n",
-    "MODEL_NAME = \"bigscience/bloomz-7b1\"  # Largest instruction-tuned BLOOM model\n",
-    "CACHE_DIR = current_dir.parent / \"data/models\"\n",
-    "CACHE_DIR.mkdir(parents=True, exist_ok=True)\n",
-    "\n",
-    "PROFESSION_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 ===\n",
-    "print(f\"Loading model: {MODEL_NAME}\")\n",
-    "print(f\"Cache directory: {CACHE_DIR}\")\n",
-    "print(f\"This may take a while on first run (~14GB download)...\\n\")\n",
-    "\n",
-    "# Check GPU availability\n",
-    "device = \"cuda\" if torch.cuda.is_available() else \"cpu\"\n",
-    "print(f\"Device: {device}\")\n",
-    "\n",
-    "if 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\n",
-    "print(\"Loading tokenizer...\")\n",
-    "try:\n",
-    "    tokenizer = AutoTokenizer.from_pretrained(\n",
-    "        MODEL_NAME,\n",
-    "        cache_dir=str(CACHE_DIR)\n",
-    "    )\n",
-    "    print(\"✅ Tokenizer loaded\")\n",
-    "except Exception as e:\n",
-    "    print(f\"❌ Error loading tokenizer: {e}\")\n",
-    "    raise\n",
-    "\n",
-    "# Ensure pad token is set\n",
-    "if tokenizer.pad_token is None:\n",
-    "    tokenizer.pad_token = tokenizer.eos_token\n",
-    "    print(f\"Set pad_token to eos_token: {tokenizer.eos_token}\")\n",
-    "\n",
-    "# Load model with optimizations\n",
-    "print(\"Loading model (this may take several minutes)...\")\n",
-    "try:\n",
-    "    model = 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",
-    "        low_cpu_mem_usage=True  # Optimize memory usage\n",
-    "    )\n",
-    "    model.eval()  # Set to evaluation mode\n",
-    "    print(\"✅ Model loaded\")\n",
-    "except Exception as e:\n",
-    "    print(f\"❌ Error loading model: {e}\")\n",
-    "    raise\n",
-    "\n",
-    "# Check VRAM usage\n",
-    "if 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 ===\n",
-    "df = pd.read_csv(input_file)\n",
-    "print(f\"Loaded {len(df)} rows\")\n",
-    "\n",
-    "if TEST_MODE:\n",
-    "    print(f\"Running in TEST MODE with {TEST_SIZE} samples\")\n",
-    "    df = df.head(TEST_SIZE).copy()\n",
-    "elif MAX_ROWS:\n",
-    "    df = df.head(MAX_ROWS).copy()\n",
-    "\n",
-    "# === CREATE PROMPT (Exact DeepSeek style) ===\n",
-    "def create_prompt(row):\n",
-    "    \"\"\"Create prompt.\"\"\"\n",
-    "    name = row.get('real_name', row.get('name', ''))\n",
-    "    if pd.isna(name):\n",
-    "        name = row.get('name', '')\n",
-    "    \n",
-    "    # Gather hints exactly like DeepSeek version\n",
-    "    hints = []\n",
-    "    if pd.notna(row.get('likely_profession')):\n",
-    "        hints.append(str(row['likely_profession']))\n",
-    "    if pd.notna(row.get('likely_nationality')):\n",
-    "        hints.append(str(row['likely_nationality']))\n",
-    "    if pd.notna(row.get('likely_country')):\n",
-    "        hints.append(str(row['likely_country']))\n",
-    "    \n",
-    "    # Add tags if we don't have enough hints\n",
-    "    if len(hints) < 3:\n",
-    "        for i in range(1, 8):\n",
-    "            tag_col = f'tag_{i}'\n",
-    "            if tag_col in row and pd.notna(row[tag_col]):\n",
-    "                tag_val = str(row[tag_col])\n",
-    "                if tag_val not in hints:\n",
-    "                    hints.append(tag_val)\n",
-    "                if len(hints) >= 5:\n",
-    "                    break\n",
-    "    \n",
-    "    hint_text = \", \".join(hints[:5]) if hints else \"none\"\n",
-    "    \n",
-    "    return f\"\"\"Given '{name}' ({hint_text}), provide:\n",
-    "1. Full legal name (Western order if non-latin script)\n",
-    "2. Any stage names/aliases (comma separated)\n",
-    "3. Gender (Male/Female/Other/Unknown)\n",
-    "4. Top 3 most likely professions from ONLY these categories:\n",
-    "   - actor\n",
-    "   - adult performer\n",
-    "   - singer/musician\n",
-    "   - model\n",
-    "   - online personality (includes streamers, cosplayers, influencers)\n",
-    "   - public figure (includes politicians, activists, journalists, authors)\n",
-    "   - voice actor/ASMR\n",
-    "   - sports professional\n",
-    "   - tv personality (includes hosts, presenters, reality TV)\n",
-    "\n",
-    "5. Primary country associated\n",
-    "\n",
-    "IMPORTANT:\n",
-    "- Choose professions ONLY from the 9 categories above\n",
-    "- Provide up to 3 professions, comma-separated, ordered by relevance\n",
-    "- Be SPECIFIC: choose the most accurate category for each role\n",
-    "- \"online personality\" includes: streamers, cosplayers, YouTubers, influencers, content creators\n",
-    "- Use 'Unknown' when uncertain or for fictional characters/places\n",
-    "- For multi-role people, list all relevant categories (e.g., \"actor, singer/musician, online personality\")\n",
-    "\n",
-    "Respond with exactly 5 numbered lines.\"\"\"\n",
-    "\n",
-    "df['prompt'] = df.apply(create_prompt, axis=1)\n",
-    "\n",
-    "# === QUERY BLOOMZ LOCAL ===\n",
-    "def query_bloomz_local(prompt: str) -> str:\n",
-    "    \"\"\"Query BLOOMZ-7B1 locally via transformers, return raw response string.\"\"\"\n",
-    "    try:\n",
-    "        # BLOOMZ works better with instruction-response format\n",
-    "        full_prompt = f\"\"\"Instruction: Extract key data on a person based on the name and hints.\n",
-    "You must respond with exactly 5 numbered lines in this format:\n",
-    "1. Full legal name\n",
-    "2. Stage names/aliases  \n",
-    "3. Gender\n",
-    "4. Professions (comma-separated, choose ONLY from: actor, adult performer, singer/musician, model, online personality, public figure, voice actor/ASMR, sports professional, tv personality)\n",
-    "5. Country\n",
-    "\n",
-    "{prompt}\n",
-    "\n",
-    "Response:\"\"\"\n",
-    "        \n",
-    "        inputs = tokenizer(\n",
-    "            full_prompt, \n",
-    "            return_tensors=\"pt\", \n",
-    "            truncation=True,\n",
-    "            max_length=2048\n",
-    "        ).to(device)\n",
-    "        \n",
-    "        # Generate with adjusted parameters for BLOOMZ\n",
-    "        with torch.no_grad():\n",
-    "            outputs = model.generate(\n",
-    "                **inputs,\n",
-    "                max_new_tokens=256,\n",
-    "                temperature=0.3,  # Increased for more variability\n",
-    "                do_sample=True,\n",
-    "                top_p=0.9,\n",
-    "                top_k=40,\n",
-    "                repetition_penalty=1.1,\n",
-    "                pad_token_id=tokenizer.eos_token_id,  # Use EOS as pad token\n",
-    "                eos_token_id=tokenizer.eos_token_id,\n",
-    "                early_stopping=True\n",
-    "            )\n",
-    "        \n",
-    "        # Decode the entire output to see what's happening\n",
-    "        full_output = tokenizer.decode(outputs[0], skip_special_tokens=True)\n",
-    "        \n",
-    "        # Extract only the generated part (after the prompt)\n",
-    "        generated_text = full_output[len(tokenizer.decode(inputs['input_ids'][0], skip_special_tokens=True)):]\n",
-    "        \n",
-    "        # Debug output\n",
-    "        if not hasattr(query_bloomz_local, 'debug_count'):\n",
-    "            query_bloomz_local.debug_count = 0\n",
-    "        \n",
-    "        if query_bloomz_local.debug_count < 3:\n",
-    "            print(f\"\\n📝 BLOOMZ Debug #{query_bloomz_local.debug_count + 1}:\")\n",
-    "            print(f\"Prompt: {full_prompt[:200]}...\")\n",
-    "            print(f\"Full output: {full_output[:500]}...\")\n",
-    "            print(f\"Generated text: {generated_text}\")\n",
-    "            print(f\"{'='*60}\\n\")\n",
-    "            query_bloomz_local.debug_count += 1\n",
-    "        \n",
-    "        return generated_text.strip()\n",
-    "        \n",
-    "    except Exception as e:\n",
-    "        print(f\"Error querying BLOOMZ: {e}\")\n",
-    "        return None\n",
-    "\n",
-    "# === PARSE RESPONSE (Exact DeepSeek format) ===\n",
-    "def parse_response(response):\n",
-    "    \"\"\"Parse numbered response into structured fields.\"\"\"\n",
-    "    if not response:\n",
-    "        return {\n",
-    "            'full_name': 'Unknown',\n",
-    "            'aliases': 'Unknown',\n",
-    "            'gender': 'Unknown',\n",
-    "            'profession_llm': 'Unknown',\n",
-    "            'country': 'Unknown'\n",
-    "        }\n",
-    "    \n",
-    "    # Split into lines and clean\n",
-    "    lines = [line.strip() for line in response.split('\\n') if line.strip()]\n",
-    "    \n",
-    "    # Initialize with Unknown values\n",
-    "    fields = {\n",
-    "        'full_name': 'Unknown',\n",
-    "        'aliases': 'Unknown',\n",
-    "        'gender': 'Unknown',\n",
-    "        'profession_llm': 'Unknown',\n",
-    "        'country': 'Unknown'\n",
-    "    }\n",
-    "    \n",
-    "    # Extract information from each numbered line\n",
-    "    for line in lines:\n",
-    "        if line.startswith('1.') or line.startswith('1)'):\n",
-    "            fields['full_name'] = line[2:].strip()\n",
-    "        elif line.startswith('2.') or line.startswith('2)'):\n",
-    "            fields['aliases'] = line[2:].strip()\n",
-    "        elif line.startswith('3.') or line.startswith('3)'):\n",
-    "            fields['gender'] = line[2:].strip()\n",
-    "        elif line.startswith('4.') or line.startswith('4)'):\n",
-    "            fields['profession_llm'] = line[2:].strip()\n",
-    "        elif line.startswith('5.') or line.startswith('5)'):\n",
-    "            fields['country'] = line[2:].strip()\n",
-    "    \n",
-    "    return fields\n",
-    "\n",
-    "# === PROCESS ===\n",
-    "output_file = current_dir.parent / f\"data/CSV/bloomz_annotated_POI{'_test' if TEST_MODE else ''}.csv\"\n",
-    "index_file = current_dir.parent / \"misc/bloomz_query_index.txt\"\n",
-    "\n",
-    "current_index = 0\n",
-    "if 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",
-    "\n",
-    "# Initialize columns (same as DeepSeek)\n",
-    "for col in ['full_name', 'gender', 'profession_llm', 'country', 'aliases']:\n",
-    "    if col not in df.columns:\n",
-    "        df[col] = 'Unknown'\n",
-    "\n",
-    "# Create prompts for all rows (same as DeepSeek)\n",
-    "print(\"Creating prompts...\")\n",
-    "df['prompt'] = df.apply(create_prompt, axis=1)\n",
-    "\n",
-    "print(f\"\\nAnnotating with BLOOMZ-7B1 LOCAL - rows {current_index} to {len(df)}...\")\n",
-    "print(f\"Model: {MODEL_NAME}\")\n",
-    "print(f\"This may take a while...\\n\")\n",
-    "\n",
-    "try:\n",
-    "    start_time = time.time()\n",
-    "    \n",
-    "    for i in tqdm(range(current_index, len(df)), desc=\"Annotating\"):\n",
-    "        row = df.iloc[i]\n",
-    "        \n",
-    "        # Query BLOOMZ (equivalent to DeepSeek query)\n",
-    "        response = query_bloomz_local(row['prompt'])\n",
-    "        parsed_data = parse_response(response)\n",
-    "        \n",
-    "        # Update dataframe\n",
-    "        for key, value in parsed_data.items():\n",
-    "            df.at[i, key] = value\n",
-    "        \n",
-    "        current_index = i + 1\n",
-    "        \n",
-    "        # Save progress at intervals\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",
-    "            print(f\"✅ Progress saved after {i+1} rows\")\n",
-    "        \n",
-    "        # Optional: Add small delay to prevent overheating (not needed for rate limiting like DeepSeek)\n",
-    "        # time.sleep(0.1)\n",
-    "    \n",
-    "    elapsed_total = time.time() - start_time\n",
-    "    print(f\"\\n✅ Done! Final results saved to {output_file}\")\n",
-    "    \n",
-    "    # Summary statistics (same as DeepSeek)\n",
-    "    print(\"\\n=== Summary Statistics ===\")\n",
-    "    print(f\"Total processed: {len(df)}\")\n",
-    "    print(f\"\\nGender distribution:\")\n",
-    "    print(df['gender'].value_counts())\n",
-    "    print(f\"\\nTop 10 profession combinations:\")\n",
-    "    print(df['profession_llm'].value_counts().head(10))\n",
-    "    print(f\"\\nTop 10 countries:\")\n",
-    "    print(df['country'].value_counts().head(10))\n",
-    "    \n",
-    "    # Sample results\n",
-    "    print(\"\\n=== Sample Results ===\")\n",
-    "    display_cols = ['real_name', 'full_name', 'gender', 'profession_llm', 'country']\n",
-    "    available_cols = [col for col in display_cols if col in df.columns]\n",
-    "    print(df[available_cols].head(10).to_string(index=False))\n",
-    "    \n",
-    "    # Additional info for local model\n",
-    "    print(f\"\\nTotal time: {elapsed_total/60:.1f} minutes\")\n",
-    "    print(f\"Average speed: {len(df)/(elapsed_total/3600):.1f} samples/hour\")\n",
-    "    if torch.cuda.is_available():\n",
-    "        print(f\"Final VRAM usage: {torch.cuda.max_memory_allocated() / 1024**3:.2f} GB\")\n",
-    "\n",
-    "except Exception as e:\n",
-    "    print(f\"⚠️ Error encountered: {e}\")\n",
-    "    print(f\"⚠️ Last processed index: {current_index}\")\n",
-    "    \n",
-    "    # Save progress before exiting\n",
-    "    df.to_csv(output_file, index=False)\n",
-    "    with open(index_file, 'w') as f:\n",
-    "        f.write(str(current_index))\n",
-    "    \n",
-    "    print(f\"⚠️ Progress saved up to row {current_index}\")"
+    "import pandas as pd",
+    "import json",
+    "import time",
+    "import re",
+    "from pathlib import Path",
+    "from tqdm import tqdm",
+    "import torch",
+    "from transformers import AutoModelForCausalLM, AutoTokenizer",
+    "from datetime import datetime",
+    "",
+    "# Import is used for pd.notna() and pd.isna() checks",
+    "",
+    "current_dir = Path.cwd())",
+    "input_file = current_dir.parent / \"data/CSV/model_adapter/real_person_adapter_step_02_NER.csv\"",
+    "",
+    "# === CONFIGURATION ===",
+    "TEST_MODE = True",
+    "TEST_SIZE = 10",
+    "MAX_ROWS = 20000",
+    "SAVE_INTERVAL = 10",
+    "",
+    "# Model settings - BLOOMZ (BigScience - European consortium)",
+    "MODEL_NAME = \"bigscience/bloomz-7b1\"  # Largest instruction-tuned BLOOM model",
+    "CACHE_DIR = current_dir.parent / \"data/models\"",
+    "CACHE_DIR.mkdir(parents=True, exist_ok=True)",
+    "",
+    "PROFESSION_CATEGORIES = [",
+    "    \"actor\", \"adult performer\", \"singer/musician\", \"model\",",
+    "    \"online personality\", \"public figure\", \"voice actor/ASMR\",",
+    "    \"sports professional\", \"tv personality\"",
+    "]",
+    "",
+    "# === LOAD MODEL ===",
+    "print(f\"Loading model: {MODEL_NAME}\")",
+    "print(f\"Cache directory: {CACHE_DIR}\")",
+    "print(f\"This may take a while on first run (~14GB download)...\\n\")",
+    "",
+    "# Check GPU availability",
+    "device = \"cuda\" if torch.cuda.is_available() else \"cpu\"",
+    "print(f\"Device: {device}\")",
+    "",
+    "if device == \"cpu\":",
+    "    print(\"\u26a0\ufe0f  WARNING: No GPU detected! Inference will be VERY slow.\")",
+    "    print(\"   Consider using a GPU or reducing model size.\")",
+    "",
+    "# Load tokenizer",
+    "print(\"Loading tokenizer...\")",
+    "try:",
+    "    tokenizer = AutoTokenizer.from_pretrained(",
+    "        MODEL_NAME,",
+    "        cache_dir=str(CACHE_DIR)",
+    "    )",
+    "    print(\"\u2705 Tokenizer loaded\")",
+    "except Exception as e:",
+    "    print(f\"\u274c Error loading tokenizer: {e}\")",
+    "    raise",
+    "",
+    "# Ensure pad token is set",
+    "if tokenizer.pad_token is None:",
+    "    tokenizer.pad_token = tokenizer.eos_token",
+    "    print(f\"Set pad_token to eos_token: {tokenizer.eos_token}\")",
+    "",
+    "# Load model with optimizations",
+    "print(\"Loading model (this may take several minutes)...\")",
+    "try:",
+    "    model = AutoModelForCausalLM.from_pretrained(",
+    "        MODEL_NAME,",
+    "        cache_dir=str(CACHE_DIR),",
+    "        torch_dtype=torch.bfloat16,  # Use BF16 for efficiency",
+    "        device_map=\"auto\",  # Automatically distribute across GPUs",
+    "        low_cpu_mem_usage=True  # Optimize memory usage",
+    "    )",
+    "    model.eval()  # Set to evaluation mode",
+    "    print(\"\u2705 Model loaded\")",
+    "except Exception as e:",
+    "    print(f\"\u274c Error loading model: {e}\")",
+    "    raise",
+    "",
+    "# Check VRAM usage",
+    "if torch.cuda.is_available():",
+    "    vram_gb = torch.cuda.max_memory_allocated() / 1024**3",
+    "    print(f\"VRAM used: {vram_gb:.2f} GB\\n\")",
+    "",
+    "# === LOAD DATA ===",
+    "df = pd.read_csv(input_file)",
+    "print(f\"Loaded {len(df)} rows\")",
+    "",
+    "if TEST_MODE:",
+    "    print(f\"Running in TEST MODE with {TEST_SIZE} samples\")",
+    "    df = df.head(TEST_SIZE).copy()",
+    "elif MAX_ROWS:",
+    "    df = df.head(MAX_ROWS).copy()",
+    "",
+    "# === CREATE PROMPT (Exact DeepSeek style) ===",
+    "def create_prompt(row):",
+    "    \"\"\"Create prompt.\"\"\"",
+    "    name = row.get('real_name', row.get('name', ''))",
+    "    if pd.isna(name):",
+    "        name = row.get('name', '')",
+    "",
+    "    # Gather hints exactly like DeepSeek version",
+    "    hints = []",
+    "    if pd.notna(row.get('likely_profession')):",
+    "        hints.append(str(row['likely_profession']))",
+    "    if pd.notna(row.get('likely_nationality')):",
+    "        hints.append(str(row['likely_nationality']))",
+    "    if pd.notna(row.get('likely_country')):",
+    "        hints.append(str(row['likely_country']))",
+    "",
+    "    # Add tags if we don't have enough hints",
+    "    if len(hints) < 3:",
+    "        for i in range(1, 8):",
+    "            tag_col = f'tag_{i}'",
+    "            if tag_col in row and pd.notna(row[tag_col]):",
+    "                tag_val = str(row[tag_col])",
+    "                if tag_val not in hints:",
+    "                    hints.append(tag_val)",
+    "                if len(hints) >= 5:",
+    "                    break",
+    "",
+    "    hint_text = \", \".join(hints[:5]) if hints else \"none\"",
+    "",
+    "    return f\"\"\"Given '{name}' ({hint_text}), provide:",
+    "1. Full legal name (Western order if non-latin script)",
+    "2. Any stage names/aliases (comma separated)",
+    "3. Gender (Male/Female/Other/Unknown)",
+    "4. Top 3 most likely professions from ONLY these categories:",
+    "   - actor",
+    "   - adult performer",
+    "   - singer/musician",
+    "   - model",
+    "   - online personality (includes streamers, cosplayers, influencers)",
+    "   - public figure (includes politicians, activists, journalists, authors)",
+    "   - voice actor/ASMR",
+    "   - sports professional",
+    "   - tv personality (includes hosts, presenters, reality TV)",
+    "",
+    "5. Primary country associated",
+    "",
+    "IMPORTANT:",
+    "- Choose professions ONLY from the 9 categories above",
+    "- Provide up to 3 professions, comma-separated, ordered by relevance",
+    "- Be SPECIFIC: choose the most accurate category for each role",
+    "- \"online personality\" includes: streamers, cosplayers, YouTubers, influencers, content creators",
+    "- Use 'Unknown' when uncertain or for fictional characters/places",
+    "- For multi-role people, list all relevant categories (e.g., \"actor, singer/musician, online personality\")",
+    "",
+    "Respond with exactly 5 numbered lines.\"\"\"",
+    "",
+    "df['prompt'] = df.apply(create_prompt, axis=1)",
+    "",
+    "# === SETUP LOGGING ===",
+    "log_dir = current_dir.parent / \"logs\"",
+    "log_dir.mkdir(parents=True, exist_ok=True)",
+    "log_file = log_dir / f\"bloomz_responses_{datetime.now().strftime('%Y%m%d_%H%M%S')}.log\"",
+    "",
+    "def log_response(idx, name, prompt, full_prompt, raw_response, parsed_data):",
+    "    \"\"\"Log all BLOOM responses to file for debugging.\"\"\"",
+    "    with open(log_file, 'a', encoding='utf-8') as f:",
+    "        f.write(f\"{'='*80}\\n\")",
+    "        f.write(f\"INDEX: {idx}\\n\")",
+    "        f.write(f\"NAME: {name}\\n\")",
+    "        f.write(f\"TIMESTAMP: {datetime.now().isoformat()}\\n\")",
+    "        f.write(f\"\\n--- ORIGINAL PROMPT ---\\n{prompt}\\n\")",
+    "        f.write(f\"\\n--- FULL PROMPT SENT TO MODEL ---\\n{full_prompt}\\n\")",
+    "        f.write(f\"\\n--- RAW MODEL RESPONSE ---\\n{raw_response}\\n\")",
+    "        f.write(f\"\\n--- PARSED DATA ---\\n{json.dumps(parsed_data, indent=2)}\\n\")",
+    "        f.write(f\"{'='*80}\\n\\n\")",
+    "",
+    "print(f\"\ud83d\udcdd Logging BLOOM responses to: {log_file}\\n\")",
+    "",
+    "# === QUERY BLOOMZ LOCAL (IMPROVED WITH FEW-SHOT) ===",
+    "def query_bloomz_local(prompt: str) -> str:",
+    "    \"\"\"Query BLOOMZ-7B1 locally via transformers, return raw response string.",
+    "",
+    "    Uses few-shot prompting which works much better for BLOOMZ.",
+    "    \"\"\"",
+    "    try:",
+    "        # Extract name and hints from the original prompt",
+    "        # Format: \"Given 'Name' (hints), provide:\"",
+    "        name_match = re.search(r\"Given '([^']+)' \\(([^)]+)\\)\", prompt)",
+    "        if name_match:",
+    "            name = name_match.group(1)",
+    "            hints = name_match.group(2)",
+    "        else:",
+    "            name = \"Unknown\"",
+    "            hints = \"no hints\"",
+    "",
+    "        # BLOOMZ responds MUCH better to few-shot examples than long instructions",
+    "        full_prompt = f\"\"\"Extract person information in 5 numbered lines.",
+    "",
+    "Example 1:",
+    "Name: Scarlett Johansson (actor, model)",
+    "1. Scarlett Ingrid Johansson",
+    "2. ScarJo",
+    "3. Female",
+    "4. actor, model",
+    "5. United States",
+    "",
+    "Example 2:",
+    "Name: Ed Sheeran (singer, musician)",
+    "1. Edward Christopher Sheeran",
+    "2. Ed Sheeran",
+    "3. Male",
+    "4. singer/musician",
+    "5. United Kingdom",
+    "",
+    "Example 3:",
+    "Name: Amouranth (streamer, cosplayer, model)",
+    "1. Kaitlyn Siragusa",
+    "2. Amouranth",
+    "3. Female",
+    "4. online personality, model",
+    "5. United States",
+    "",
+    "Now extract for:",
+    "Name: {name} ({hints})",
+    "1.\"\"\"",
+    "",
+    "        # Store full prompt for logging",
+    "        query_bloomz_local.last_full_prompt = full_prompt",
+    "",
+    "        inputs = tokenizer(",
+    "            full_prompt,",
+    "            return_tensors=\"pt\",",
+    "            truncation=True,",
+    "            max_length=2048",
+    "        ).to(device)",
+    "",
+    "        # Generate with adjusted parameters for BLOOMZ",
+    "        with torch.no_grad():",
+    "            outputs = model.generate(",
+    "                **inputs,",
+    "                max_new_tokens=256,",
+    "                temperature=0.3,  # Increased for more variability",
+    "                do_sample=True,",
+    "                top_p=0.9,",
+    "                top_k=40,",
+    "                repetition_penalty=1.1,",
+    "                pad_token_id=tokenizer.eos_token_id,  # Use EOS as pad token",
+    "                eos_token_id=tokenizer.eos_token_id,",
+    "                early_stopping=True",
+    "            )",
+    "",
+    "        # Decode the entire output to see what's happening",
+    "        full_output = tokenizer.decode(outputs[0], skip_special_tokens=True)",
+    "",
+    "        # Extract only the generated part (after the prompt)",
+    "        generated_text = full_output[len(tokenizer.decode(inputs['input_ids'][0], skip_special_tokens=True)):]",
+    "",
+    "        # Debug output (show first 3)",
+    "        if not hasattr(query_bloomz_local, 'debug_count'):",
+    "            query_bloomz_local.debug_count = 0",
+    "",
+    "        if query_bloomz_local.debug_count < 3:",
+    "            print(f\"\\n\ud83d\udcdd BLOOMZ Debug #{query_bloomz_local.debug_count + 1}:\")",
+    "            print(f\"Input: {name} ({hints})\")",
+    "            print(f\"Generated: {generated_text[:200]}\")",
+    "            print(f\"{'='*60}\\n\")",
+    "            query_bloomz_local.debug_count += 1",
+    "",
+    "        return generated_text.strip()",
+    "",
+    "    except Exception as e:",
+    "        print(f\"Error querying BLOOMZ: {e}\")",
+    "        query_bloomz_local.last_full_prompt = f\"ERROR: {e}\"",
+    "        return None",
+    "",
+    "# === PARSE RESPONSE (Exact DeepSeek format) ===",
+    "def parse_response(response):",
+    "    \"\"\"Parse numbered response into structured fields.\"\"\"",
+    "    if not response:",
+    "        return {",
+    "            'full_name': 'Unknown',",
+    "            'aliases': 'Unknown',",
+    "            'gender': 'Unknown',",
+    "            'profession_llm': 'Unknown',",
+    "            'country': 'Unknown'",
+    "        }",
+    "",
+    "    # Split into lines and clean",
+    "    lines = [line.strip() for line in response.split('\\n') if line.strip()]",
+    "",
+    "    # Initialize with Unknown values",
+    "    fields = {",
+    "        'full_name': 'Unknown',",
+    "        'aliases': 'Unknown',",
+    "        'gender': 'Unknown',",
+    "        'profession_llm': 'Unknown',",
+    "        'country': 'Unknown'",
+    "    }",
+    "",
+    "    # Extract information from each numbered line",
+    "    for line in lines:",
+    "        if line.startswith('1.') or line.startswith('1)'):",
+    "            fields['full_name'] = line[2:].strip()",
+    "        elif line.startswith('2.') or line.startswith('2)'):",
+    "            fields['aliases'] = line[2:].strip()",
+    "        elif line.startswith('3.') or line.startswith('3)'):",
+    "            fields['gender'] = line[2:].strip()",
+    "        elif line.startswith('4.') or line.startswith('4)'):",
+    "            fields['profession_llm'] = line[2:].strip()",
+    "        elif line.startswith('5.') or line.startswith('5)'):",
+    "            fields['country'] = line[2:].strip()",
+    "",
+    "    return fields",
+    "",
+    "# === PROCESS ===",
+    "output_file = current_dir.parent / f\"data/CSV/bloomz_annotated_POI{'_test' if TEST_MODE else ''}.csv\"",
+    "index_file = current_dir.parent / \"misc/bloomz_query_index.txt\"",
+    "",
+    "current_index = 0",
+    "if index_file.exists():",
+    "    with open(index_file) as f:",
+    "        current_index = int(f.read().strip())",
+    "    print(f\"Resuming from index {current_index}\")",
+    "",
+    "# Initialize columns (same as DeepSeek)",
+    "for col in ['full_name', 'gender', 'profession_llm', 'country', 'aliases']:",
+    "    if col not in df.columns:",
+    "        df[col] = 'Unknown'",
+    "",
+    "# Create prompts for all rows (same as DeepSeek)",
+    "print(\"Creating prompts...\")",
+    "df['prompt'] = df.apply(create_prompt, axis=1)",
+    "",
+    "print(f\"\\nAnnotating with BLOOMZ-7B1 LOCAL - rows {current_index} to {len(df)}...\")",
+    "print(f\"Model: {MODEL_NAME}\")",
+    "print(f\"This may take a while...\\n\")",
+    "",
+    "try:",
+    "    start_time = time.time()",
+    "",
+    "    for i in tqdm(range(current_index, len(df)), desc=\"Annotating\"):",
+    "        row = df.iloc[i]",
+    "",
+    "        # Query BLOOMZ (equivalent to DeepSeek query)",
+    "        response = query_bloomz_local(row['prompt'])",
+    "        parsed_data = parse_response(response)",
+    "",
+    "        # Log the complete interaction for debugging",
+    "        log_response(",
+    "            idx=i,",
+    "            name=row.get('real_name', row.get('name', 'Unknown')),",
+    "            prompt=row['prompt'],",
+    "            full_prompt=query_bloomz_local.last_full_prompt if hasattr(query_bloomz_local, 'last_full_prompt') else 'N/A',",
+    "            raw_response=response if response else 'None',",
+    "            parsed_data=parsed_data",
+    "        )",
+    "",
+    "        # Update dataframe",
+    "        for key, value in parsed_data.items():",
+    "            df.at[i, key] = value",
+    "",
+    "        current_index = i + 1",
+    "",
+    "        # Save progress at intervals",
+    "        if (i + 1) % SAVE_INTERVAL == 0 or (i + 1) == len(df):",
+    "            df.to_csv(output_file, index=False)",
+    "            with open(index_file, 'w') as f:",
+    "                f.write(str(current_index))",
+    "            print(f\"\u2705 Progress saved after {i+1} rows\")",
+    "",
+    "        # Optional: Add small delay to prevent overheating (not needed for rate limiting like DeepSeek)",
+    "        # time.sleep(0.1)",
+    "",
+    "    elapsed_total = time.time() - start_time",
+    "    print(f\"\\n\u2705 Done! Final results saved to {output_file}\")",
+    "",
+    "    # Summary statistics (same as DeepSeek)",
+    "    print(\"\\n=== Summary Statistics ===\")",
+    "    print(f\"Total processed: {len(df)}\")",
+    "    print(f\"\\nGender distribution:\")",
+    "    print(df['gender'].value_counts())",
+    "    print(f\"\\nTop 10 profession combinations:\")",
+    "    print(df['profession_llm'].value_counts().head(10))",
+    "    print(f\"\\nTop 10 countries:\")",
+    "    print(df['country'].value_counts().head(10))",
+    "",
+    "    # Sample results",
+    "    print(\"\\n=== Sample Results ===\")",
+    "    display_cols = ['real_name', 'full_name', 'gender', 'profession_llm', 'country']",
+    "    available_cols = [col for col in display_cols if col in df.columns]",
+    "    print(df[available_cols].head(10).to_string(index=False))",
+    "",
+    "    # Additional info for local model",
+    "    print(f\"\\nTotal time: {elapsed_total/60:.1f} minutes\")",
+    "    print(f\"Average speed: {len(df)/(elapsed_total/3600):.1f} samples/hour\")",
+    "    if torch.cuda.is_available():",
+    "        print(f\"Final VRAM usage: {torch.cuda.max_memory_allocated() / 1024**3:.2f} GB\")",
+    "",
+    "except Exception as e:",
+    "    print(f\"\u26a0\ufe0f Error encountered: {e}\")",
+    "    print(f\"\u26a0\ufe0f Last processed index: {current_index}\")",
+    "",
+    "    # Save progress before exiting",
+    "    df.to_csv(output_file, index=False)",
+    "    with open(index_file, 'w') as f:",
+    "        f.write(str(current_index))",
+    "",
+    "    print(f\"\u26a0\ufe0f Progress saved up to row {current_index}\")",
+    ""
    ]
   }
  ],
@@ -367,4 +433,4 @@
  },
  "nbformat": 4,
  "nbformat_minor": 5
-}
+}