Skip to content

Conversation

Zeyi-Lin
Copy link
Member

@Zeyi-Lin Zeyi-Lin commented Apr 19, 2024

Description

与huggingface-transformers库的集成,测试代码为:

import evaluate
import numpy as np
import swanlab
from swanlab.integration.huggingface import SwanLabCallback
from datasets import load_dataset
from transformers import AutoModelForSequenceClassification, AutoTokenizer, Trainer, TrainingArguments


def tokenize_function(examples):
    return tokenizer(examples["text"], padding="max_length", truncation=True)


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


dataset = load_dataset("yelp_review_full")

tokenizer = AutoTokenizer.from_pretrained("bert-base-cased")

tokenized_datasets = dataset.map(tokenize_function, batched=True)

small_train_dataset = tokenized_datasets["train"].shuffle(seed=42).select(range(1000))
small_eval_dataset = tokenized_datasets["test"].shuffle(seed=42).select(range(1000))

metric = evaluate.load("accuracy")

model = AutoModelForSequenceClassification.from_pretrained("bert-base-cased", num_labels=5)

training_args = TrainingArguments(
    output_dir="test_trainer",
    report_to="none",
    num_train_epochs=3,
    logging_steps=50,
)

swanlab_callback = SwanLabCallback(experiment_name="TransformersTest")


trainer = Trainer(
    model=model,
    args=training_args,
    train_dataset=small_train_dataset,
    eval_dataset=small_eval_dataset,
    compute_metrics=compute_metrics,
    callbacks=[swanlab_callback],
)

trainer.train()

Closes: #481

@Zeyi-Lin Zeyi-Lin requested a review from SAKURA-CAT April 19, 2024 08:40
@Zeyi-Lin Zeyi-Lin self-assigned this Apr 19, 2024
@Zeyi-Lin Zeyi-Lin added this to the Integration milestone Apr 19, 2024
@Zeyi-Lin Zeyi-Lin added the 💪 enhancement New feature or request label Apr 19, 2024
@Zeyi-Lin Zeyi-Lin requested a review from SAKURA-CAT April 28, 2024 06:47
@SAKURA-CAT SAKURA-CAT merged commit af59c36 into main Apr 28, 2024
@SAKURA-CAT SAKURA-CAT deleted the feat/integration-huggingface-transformers branch April 28, 2024 07:49
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Labels
💪 enhancement New feature or request
Projects
None yet
Development

Successfully merging this pull request may close these issues.

[REQUEST] 集成HuggingFace Transformers
2 participants