HuggingFace, PyTorch, 그리고 Weights & Biases(W&B)로 구축하는 개체명 인식 시스템

!pip install seqeval
!pip install transformers
!pip install datasets

import torch
from transformers import AutoTokenizer, AutoModelForTokenClassification, Trainer, TrainingArguments
from datasets import load_dataset, load_metric, Dataset, DatasetDict
import numpy as np
from seqeval.metrics import f1_score, precision_score, recall_score, classification_report

print("CUDA available:", torch.cuda.is_available())
print("Current device index:", torch.cuda.current_device())
print("Device name:", torch.cuda.get_device_name(torch.cuda.current_device()))

def read_conll_file(file_path):
    with open(file_path, "r") as f:
        content = f.read().strip()
        sentences = content.split("\n\n")
        data = []
        for sentence in sentences:
            tokens = sentence.split("\n")
            token_data = []
            for token in tokens:
                token_data.append(token.split())
            data.append(token_data)
    return data

train_data = read_conll_file("/kaggle/input/conll2003-dataset/conll2003/eng.train")
validation_data = read_conll_file("/kaggle/input/conll2003-dataset/conll2003/eng.testa")
test_data = read_conll_file("/kaggle/input/conll2003-dataset/conll2003/eng.testb")

def convert_to_dataset(data, label_map):
    formatted_data = {"tokens": [], "ner_tags": []}
    for sentence in data:
        tokens = [token_data[0] for token_data in sentence]
        ner_tags = [label_map[token_data[3]] for token_data in sentence]
        formatted_data["tokens"].append(tokens)
        formatted_data["ner_tags"].append(ner_tags)
    return Dataset.from_dict(formatted_data)

label_list = sorted(list(set([token_data[3] for sentence in train_data for token_data in sentence])))
label_map = {label: i for i, label in enumerate(label_list)}

train_dataset = convert_to_dataset(train_data, label_map)
validation_dataset = convert_to_dataset(validation_data, label_map)
test_dataset = convert_to_dataset(test_data, label_map)

datasets = DatasetDict({
    "train": train_dataset,
    "validation": validation_dataset,
    "test": test_dataset,
})

model_name = "bert-base-cased"
tokenizer = AutoTokenizer.from_pretrained(model_name)
model = AutoModelForTokenClassification.from_pretrained(model_name, num_labels=len(label_list))

def compute_metrics(eval_prediction):
    predictions, labels = eval_prediction
    predictions = np.argmax(predictions, axis=2)

    # Remove ignored index (special tokens)
    true_predictions = [
        [label_list[p] for (p, l) in zip(prediction, label) if l != -100]
        for prediction, label in zip(predictions, labels)
    ]
    true_labels = [
        [label_list[l] for (p, l) in zip(prediction, label) if l != -100]
        for prediction, label in zip(predictions, labels)
    ]

    return {
        "precision": precision_score(true_labels, true_predictions),
        "recall": recall_score(true_labels, true_predictions),
        "f1": f1_score(true_labels, true_predictions),
        "classification_report": classification_report(true_labels, true_predictions),
    }

def tokenize_and_align_labels(examples):
    tokenized_inputs = tokenizer(
        examples["tokens"], truncation=True, is_split_into_words=True, padding=True
    )
    labels = []
    for i, label in enumerate(examples["ner_tags"]):
        word_ids = tokenized_inputs.word_ids(batch_index=i)
        previous_word_idx = None
        label_ids = []
        for word_idx in word_ids:
            if word_idx is None:
                label_ids.append(-100)
            elif word_idx != previous_word_idx:
                label_ids.append(label[word_idx])
            else:
                label_ids.append(-100)
            previous_word_idx = word_idx
        labels.append(label_ids)
    tokenized_inputs["labels"] = labels
    return tokenized_inputs

tokenized_datasets = datasets.map(tokenize_and_align_labels, batched=True)

training_args = TrainingArguments(
    output_dir="./results",
    evaluation_strategy="steps",
    eval_steps=500,
    save_steps=500,
    num_train_epochs=1,
    per_device_train_batch_size=8,
    per_device_eval_batch_size=8,
    logging_steps=100,
    learning_rate=5e-5,
    load_best_model_at_end=True,
    metric_for_best_model="f1",
)

def data_collator(data):
    input_ids = [torch.tensor(item["input_ids"]) for item in data]
    attention_mask = [torch.tensor(item["attention_mask"]) for item in data]
    labels = [torch.tensor(item["labels"]) for item in data]

    input_ids = torch.nn.utils.rnn.pad_sequence(input_ids, batch_first=True, padding_value=tokenizer.pad_token_id)
    attention_mask = torch.nn.utils.rnn.pad_sequence(attention_mask, batch_first=True, padding_value=0)
    labels = torch.nn.utils.rnn.pad_sequence(labels, batch_first=True, padding_value=-100)

    return {
        "input_ids": input_ids,
        "attention_mask": attention_mask,
        "labels": labels,
    }

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

trainer.train()

sentence = "John Smith is a software engineer who works at Google."

tokenized_input = tokenizer(sentence, return_tensors="pt").to(model.device)

outputs = model(**tokenized_input)

predicted_labels = outputs.logits.argmax(-1)[0]

named_entities = [tokenizer.decode([token]) for token, label in zip(tokenized_input["input_ids"][0], predicted_labels) if label != 0 and label != label_map['O']]

print("Named Entities - Example 1:", named_entities)

sentence2 = "The company Apple Inc. announced its new product, the iPhone 12, at a press conference held in San Francisco."

tokenized_input2 = tokenizer(sentence2, return_tensors="pt").to(model.device)

outputs2 = model(**tokenized_input2)

predicted_labels2 = outputs2.logits.argmax(-1)[0]

named_entities2 = [tokenizer.decode([token]) for token, label in zip(tokenized_input2["input_ids"][0], predicted_labels2) if label != 0 and label != label_map['O']]

print("Named Entities - Example 2:", named_entities2)

sentence3 = "The actor Tom Hanks starred in the movie Forrest Gump."

tokenized_input3 = tokenizer(sentence3, return_tensors="pt").to(model.device)

outputs3 = model(**tokenized_input3)

predicted_labels3 = outputs3.logits.argmax(-1)[0]

named_entities3 = [tokenizer.decode([token]) for token, label in zip(tokenized_input3["input_ids"][0], predicted_labels3) if label != 0 and label != label_map['O']]

print("Named Entities - Example 3:", named_entities3)

sentence4 = "Paris is the capital city of France."

tokenized_input4 = tokenizer(sentence4, return_tensors="pt").to(model.device)

outputs4 = model(**tokenized_input4)

predicted_labels4 = outputs4.logits.argmax(-1)[0]

named_entities4 = [tokenizer.decode([token]) for token, label in zip(tokenized_input4["input_ids"][0], predicted_labels4) if label != 0 and label != label_map['O']]

print("Named Entities - Example 4:", named_entities4)

sentence5 = "The scientist Marie Curie won the Nobel Prize in Physics and Chemistry."

tokenized_input5 = tokenizer(sentence5, return_tensors="pt").to(model.device)

outputs5 = model(**tokenized_input5)

predicted_labels5 = outputs5.logits.argmax(-1)[0]

named_entities5 = [tokenizer.decode([token]) for token, label in zip(tokenized_input5["input_ids"][0], predicted_labels5) if label != 0 and label != label_map['O']]

print("Named Entities - Example 5:", named_entities5)