from transformers import pipeline

# ===========================
# 1. NER pipeline — instant, zero training!
# ===========================
ner = pipeline("ner", grouped_entities=True, device=0)

text = "Joko Widodo met Tim Cook at Apple Park in Cupertino on January 15, 2024."
entities = ner(text)

for e in entities:
    print(f"  {e['word']:20s} → {e['entity_group']:5s} ({e['score']:.1%})  [{e['start']}:{e['end']}]")
# Joko Widodo          → PER   (99.8%)  [0:12]
# Tim Cook             → PER   (99.6%)  [17:25]
# Apple Park           → LOC   (97.3%)  [29:39]
# Cupertino            → LOC   (99.9%)  [43:52]
# January 15, 2024     → MISC  (85.2%)  [56:73]

# ===========================
# 2. Entity types (standard CoNLL-2003)
# ===========================
# PER  = Person         (Joko Widodo, Tim Cook, Albert Einstein)
# LOC  = Location       (Jakarta, California, Mount Everest)
# ORG  = Organization   (Google, BRI, United Nations)
# MISC = Miscellaneous  (Indonesian, COVID-19, iPhone 15)
#
# Custom NER bisa menambah tipe apapun:
# PRODUCT, DATE, MONEY, EVENT, DISEASE, DRUG, dll.

# ===========================
# 3. NER use cases di production
# ===========================
# • Search engines: extract entities for knowledge graphs
# • Customer support: detect product names, order IDs
# • Medical NLP: extract drug names, diseases, symptoms
# • Finance: extract company names, monetary amounts
# • Legal: extract person names, dates, legal references
# • Social media: extract locations, organizations, events

# ===========================
# Standard CoNLL-2003 label list (9 labels)
# ===========================
label_list = [
    "O",       # 0 — Outside (bukan entitas)
    "B-PER",   # 1 — Beginning of Person
    "I-PER",   # 2 — Inside Person
    "B-ORG",   # 3 — Beginning of Organization
    "I-ORG",   # 4 — Inside Organization
    "B-LOC",   # 5 — Beginning of Location
    "I-LOC",   # 6 — Inside Location
    "B-MISC",  # 7 — Beginning of Miscellaneous
    "I-MISC",  # 8 — Inside Miscellaneous
]

# Mappings
label2id = {label: i for i, label in enumerate(label_list)}
id2label = {i: label for i, label in enumerate(label_list)}

print(label2id)
# {'O': 0, 'B-PER': 1, 'I-PER': 2, 'B-ORG': 3, 'I-ORG': 4,
#  'B-LOC': 5, 'I-LOC': 6, 'B-MISC': 7, 'I-MISC': 8}

# num_labels untuk model = 9
# Formula: num_labels = 1 (O) + 2 × num_entity_types (B + I per type)
# 4 entity types → 1 + 2×4 = 9 labels

from datasets import load_dataset

# ===========================
# Load CoNLL-2003
# ===========================
dataset = load_dataset("conll2003")
print(dataset)
# DatasetDict({
#     train: Dataset({features: ['id','tokens','pos_tags','chunk_tags','ner_tags'], num_rows: 14041})
#     validation: Dataset({num_rows: 3250})
#     test: Dataset({num_rows: 3453})
# })

# ===========================
# Inspect satu contoh
# ===========================
example = dataset["train"][0]
print(f"Tokens:   {example['tokens']}")
print(f"NER tags: {example['ner_tags']}")
# Tokens:   ['EU', 'rejects', 'German', 'call', 'to', 'boycott', 'British', 'lamb', '.']
# NER tags: [3, 0, 7, 0, 0, 0, 7, 0, 0]
#           B-ORG O  B-MISC O O O  B-MISC O O

# Decode NER tags to human-readable
ner_feature = dataset["train"].features["ner_tags"].feature
print(f"Label names: {ner_feature.names}")
# ['O', 'B-PER', 'I-PER', 'B-ORG', 'I-ORG', 'B-LOC', 'I-LOC', 'B-MISC', 'I-MISC']

# Human-readable version:
for token, tag_id in zip(example["tokens"], example["ner_tags"]):
    tag = ner_feature.names[tag_id]
    marker = "  ←" if tag != "O" else ""
    print(f"  {token:15s} {tag:8s}{marker}")
# EU              B-ORG   ←   (European Union)
# rejects         O
# German          B-MISC  ←   (nationality)
# call            O
# to              O
# boycott         O
# British         B-MISC  ←   (nationality)
# lamb            O
# .               O

# PENTING: Data sudah PRE-TOKENIZED (list of words + list of tags)
# Ini BERBEDA dari text classification di mana input = string!
# Input: ["EU", "rejects", "German", ...] (list, bukan string!)
# Labels: [3, 0, 7, ...] (satu label per kata)

from transformers import AutoTokenizer

tokenizer = AutoTokenizer.from_pretrained("bert-base-cased")
# IMPORTANT: use CASED model for NER! Names are capitalized!

# ===========================
# THE alignment function — THE MOST IMPORTANT CODE IN NER!
# ===========================
def tokenize_and_align_labels(examples):
    """Tokenize pre-tokenized words and align NER labels.

    Key concept: word_ids() tells us which WORD each subword came from.
    - word_id=None → special token ([CLS], [SEP], [PAD]) → label=-100
    - First subword of a word → keep original label
    - Continuation subword → label=-100

    Args:
        examples: batch with 'tokens' (list of words) and 'ner_tags' (list of ints)
    Returns:
        dict with input_ids, attention_mask, labels (aligned!)
    """
    tokenized = tokenizer(
        examples["tokens"],
        truncation=True,
        is_split_into_words=True,  # ← CRITICAL! Input is already split into words!
        max_length=128,
    )

    all_labels = []
    for i, labels in enumerate(examples["ner_tags"]):
        word_ids = tokenized.word_ids(batch_index=i)
        # word_ids: [None, 0, 0, 1, 1, 1, 2, 3, 3, 3, None]
        #           [CLS] Jo ##ko Wi ##do ##do visited Cup ##ert ##ino [SEP]

        label_ids = []
        previous_word_id = None

        for word_id in word_ids:
            if word_id is None:
                # Special token → ignore
                label_ids.append(-100)
            elif word_id != previous_word_id:
                # First subword of a new word → use word's label
                label_ids.append(labels[word_id])
            else:
                # Continuation subword → ignore
                label_ids.append(-100)

            previous_word_id = word_id

        all_labels.append(label_ids)

    tokenized["labels"] = all_labels
    return tokenized

# ===========================
# Apply to dataset
# ===========================
tokenized_dataset = dataset.map(
    tokenize_and_align_labels,
    batched=True,
    remove_columns=dataset["train"].column_names,
)

# Verify alignment
example = tokenized_dataset["train"][0]
tokens = tokenizer.convert_ids_to_tokens(example["input_ids"])
labels = example["labels"]
for token, label in zip(tokens, labels):
    label_name = label_list[label] if label != -100 else "IGNORE"
    print(f"  {token:15s} → {label_name}")
# [CLS]           → IGNORE
# EU              → B-ORG   ← first subword, keep label
# re              → IGNORE  ← continuation, ignore
# ##ject          → IGNORE
# ##s             → IGNORE
# German          → B-MISC  ← first subword of "German"
# ...

import numpy as np
import evaluate
from datasets import load_dataset
from transformers import (
    AutoTokenizer, AutoModelForTokenClassification,
    TrainingArguments, Trainer, DataCollatorForTokenClassification
)

# ═══════════════════════════════════════
# STEP 1: LOAD DATASET
# ═══════════════════════════════════════
dataset = load_dataset("conll2003")
label_list = dataset["train"].features["ner_tags"].feature.names
num_labels = len(label_list)  # 9
print(f"Labels ({num_labels}): {label_list}")

# ═══════════════════════════════════════
# STEP 2: LOAD TOKENIZER & MODEL
# ═══════════════════════════════════════
model_name = "bert-base-cased"  # CASED for NER!
tokenizer = AutoTokenizer.from_pretrained(model_name)
model = AutoModelForTokenClassification.from_pretrained(
    model_name,
    num_labels=num_labels,
    id2label={i: l for i, l in enumerate(label_list)},
    label2id={l: i for i, l in enumerate(label_list)},
)

# ═══════════════════════════════════════
# STEP 3: TOKENIZE + ALIGN (from Section 6!)
# ═══════════════════════════════════════
def tokenize_and_align(examples):
    tokenized = tokenizer(examples["tokens"], truncation=True,
                          is_split_into_words=True, max_length=128)
    all_labels = []
    for i, labels in enumerate(examples["ner_tags"]):
        word_ids = tokenized.word_ids(batch_index=i)
        label_ids, prev = [], None
        for wid in word_ids:
            if wid is None:       label_ids.append(-100)
            elif wid != prev:     label_ids.append(labels[wid])
            else:                 label_ids.append(-100)
            prev = wid
        all_labels.append(label_ids)
    tokenized["labels"] = all_labels
    return tokenized

tokenized = dataset.map(tokenize_and_align, batched=True,
                        remove_columns=dataset["train"].column_names)

# ═══════════════════════════════════════
# STEP 4: DATA COLLATOR (special for token classification!)
# ═══════════════════════════════════════
data_collator = DataCollatorForTokenClassification(tokenizer=tokenizer)
# Pads BOTH input_ids AND labels to same length!
# Padding labels get value -100 (ignored in loss)

# ═══════════════════════════════════════
# STEP 5: METRICS — seqeval (NER-specific!)
# ═══════════════════════════════════════
seqeval = evaluate.load("seqeval")

def compute_metrics(eval_pred):
    logits, labels = eval_pred
    predictions = np.argmax(logits, axis=-1)

    # Convert IDs back to label strings (skip -100!)
    true_labels, true_preds = [], []
    for pred_seq, label_seq in zip(predictions, labels):
        t_labels, t_preds = [], []
        for p, l in zip(pred_seq, label_seq):
            if l != -100:  # skip ignored positions!
                t_labels.append(label_list[l])
                t_preds.append(label_list[p])
        true_labels.append(t_labels)
        true_preds.append(t_preds)

    results = seqeval.compute(predictions=true_preds, references=true_labels)
    return {
        "precision": results["overall_precision"],
        "recall": results["overall_recall"],
        "f1": results["overall_f1"],
        "accuracy": results["overall_accuracy"],
    }

# ═══════════════════════════════════════
# STEP 6: TRAIN!
# ═══════════════════════════════════════
args = TrainingArguments(
    output_dir="./ner-bert",
    num_train_epochs=3,
    per_device_train_batch_size=16,
    per_device_eval_batch_size=64,
    learning_rate=2e-5,
    weight_decay=0.01,
    warmup_ratio=0.1,
    fp16=True,
    eval_strategy="epoch",
    save_strategy="epoch",
    load_best_model_at_end=True,
    metric_for_best_model="f1",
    save_total_limit=2,
    report_to="none",
)

trainer = Trainer(
    model=model, args=args,
    train_dataset=tokenized["train"],
    eval_dataset=tokenized["validation"],
    tokenizer=tokenizer,
    data_collator=data_collator,
    compute_metrics=compute_metrics,
)

trainer.train()

# ═══════════════════════════════════════
# STEP 7: EVALUATE
# ═══════════════════════════════════════
results = trainer.evaluate(tokenized["test"])
print(f"\n🏷️ NER Results:")
print(f"  F1:        {results['eval_f1']:.1%}")
print(f"  Precision: {results['eval_precision']:.1%}")
print(f"  Recall:    {results['eval_recall']:.1%}")
# 🏷️ NER Results:
#   F1:        91.8%
#   Precision: 92.1%
#   Recall:    91.5%

# Save
trainer.save_model("./ner-bert-final")
tokenizer.save_pretrained("./ner-bert-final")
print("🏆 NER model saved!")

import evaluate

seqeval = evaluate.load("seqeval")

# Example predictions and references
true_labels = [["B-PER", "I-PER", "O", "B-LOC"]]
pred_labels = [["B-PER", "I-PER", "O", "B-ORG"]]  # LOC→ORG mistake!

results = seqeval.compute(predictions=pred_labels, references=true_labels)
print(results)
# {
#   'PER':  {'precision': 1.0,  'recall': 1.0,  'f1': 1.0,  'number': 1},
#   'LOC':  {'precision': 0.0,  'recall': 0.0,  'f1': 0.0,  'number': 1},  ← missed!
#   'ORG':  {'precision': 0.0,  'recall': 0.0,  'f1': 0.0,  'number': 0},  ← false positive
#   'overall_precision': 0.5,
#   'overall_recall': 0.5,
#   'overall_f1': 0.5,
#   'overall_accuracy': 0.75,
# }

# PENTING: seqeval mengevaluasi pada level ENTITAS, bukan token!
# "Joko Widodo" = 1 entitas PER (2 tokens)
# Jika model memprediksi "Joko"=B-PER tapi "Widodo"=O (bukan I-PER)
# → entitas SALAH (partial match = wrong!)
# seqeval menggunakan "exact match" — semua tokens harus benar.

# Typical results after fine-tuning BERT on CoNLL-2003:
# PER:  F1 = 96.2% (names are relatively easy)
# LOC:  F1 = 93.1% (locations are clear)
# ORG:  F1 = 89.4% (organizations can be ambiguous)
# MISC: F1 = 82.6% (misc is hardest — very diverse)

from transformers import pipeline

# ===========================
# 1. Use fine-tuned model as pipeline (EASIEST!)
# ===========================
ner_pipe = pipeline("ner",
    model="./ner-bert-final",
    aggregation_strategy="simple",  # merge subwords automatically!
    device=0,
)
# aggregation_strategy options:
# "none"   → raw subword predictions (no merging)
# "simple" → merge subwords, take first subword's label ← RECOMMENDED
# "first"  → same as simple
# "average" → average scores across subwords
# "max"    → take max score across subwords

text = "Barack Obama visited the United Nations headquarters in New York City."
entities = ner_pipe(text)
for e in entities:
    print(f"  {e['word']:25s} {e['entity_group']:5s} ({e['score']:.1%})")
# Barack Obama              PER   (99.7%)
# United Nations            ORG   (99.2%)
# New York City             LOC   (99.5%)

# ===========================
# 2. Manual post-processing (for custom needs)
# ===========================
import torch

def extract_entities(text, model, tokenizer):
    """Extract entities manually with full control."""
    inputs = tokenizer(text, return_tensors="pt", return_offsets_mapping=True)
    offsets = inputs.pop("offset_mapping")[0]

    with torch.no_grad():
        outputs = model(**inputs.to(model.device))
    predictions = outputs.logits.argmax(dim=-1)[0].cpu().tolist()

    entities = []
    current_entity = None

    for idx, (pred, offset) in enumerate(zip(predictions, offsets)):
        label = model.config.id2label[pred]
        start, end = offset.tolist()

        if start == 0 and end == 0:  # special token
            continue

        if label.startswith("B-"):
            if current_entity:
                entities.append(current_entity)
            current_entity = {
                "entity": label[2:],
                "word": text[start:end],
                "start": start, "end": end
            }
        elif label.startswith("I-") and current_entity:
            current_entity["word"] = text[current_entity["start"]:end]
            current_entity["end"] = end
        else:
            if current_entity:
                entities.append(current_entity)
                current_entity = None

    if current_entity:
        entities.append(current_entity)
    return entities

result = extract_entities(text, model, tokenizer)
for e in result:
    print(f"  {e['word']:25s} → {e['entity']}")

from datasets import Dataset, Features, Sequence, ClassLabel, Value

# ===========================
# Format 1: CoNLL-style (recommended)
# ===========================
# File: data.txt (tab-separated, empty line = sentence boundary)
# Joko    B-PER
# Widodo  I-PER
# visited O
# BRI     B-ORG
# .       O
#                  ← empty line = new sentence
# Jakarta B-LOC
# is      O
# big     O
# .       O

# ===========================
# Format 2: From Python dict (easiest for small datasets)
# ===========================
custom_labels = ["O", "B-PER", "I-PER", "B-ORG", "I-ORG", "B-LOC", "I-LOC"]

data = {
    "tokens": [
        ["Joko", "Widodo", "visited", "BRI", "in", "Jakarta"],
        ["Google", "opened", "office", "in", "Surabaya"],
    ],
    "ner_tags": [
        [1, 2, 0, 3, 0, 5],   # B-PER I-PER O B-ORG O B-LOC
        [3, 0, 0, 0, 5],       # B-ORG O O O B-LOC
    ],
}

features = Features({
    "tokens": Sequence(Value("string")),
    "ner_tags": Sequence(ClassLabel(names=custom_labels)),
})

custom_dataset = Dataset.from_dict(data, features=features)
split = custom_dataset.train_test_split(test_size=0.2, seed=42)

# Now use EXACTLY the same pipeline as Section 8!
# Just replace: dataset = split  (instead of load_dataset("conll2003"))
# And update: label_list = custom_labels
# Everything else is IDENTICAL. 🎉

# Minimum data recommendation:
# • 200-500 annotated sentences for basic NER
# • 1000-5000 for production-quality
# • 10000+ for state-of-the-art

# POS tags in CoNLL-2003:
# "EU rejects German call to boycott British lamb ."
# NNP VBZ    JJ     NN   TO VB      JJ      NN   .
# (Proper Noun, Verb, Adjective, Noun, To, Verb, Adj, Noun, Punct)

# Fine-tuning is IDENTICAL to NER!
# Just change:
# 1. label_list → POS tag list (NN, VB, JJ, DT, ...)
# 2. examples["ner_tags"] → examples["pos_tags"]
# 3. num_labels → len(pos_tag_list)
# Everything else (alignment, collator, trainer) is the SAME!

# Other token classification tasks:
# • Chunking (NP, VP, PP phrases)
# • Aspect-Based Sentiment (target + sentiment per word)
# • Keyphrase extraction
# • Slot filling (for chatbot intents)