Transformer / app.py
flytoe's picture
Update app.py
849f804 verified
raw
history blame
1.53 kB
import torch
from transformers import AutoModelForSequenceClassification, Trainer, TrainingArguments, AutoTokenizer
from datasets import load_dataset
# 1️⃣ Modell & Tokenizer laden
model_name = "allenai/scibert_scivocab_uncased"
tokenizer = AutoTokenizer.from_pretrained(model_name)
model = AutoModelForSequenceClassification.from_pretrained(model_name, num_labels=3)
# 2️⃣ Dataset laden (armanc/scientific_papers) mit trust_remote_code=True
dataset = load_dataset("armanc/scientific_papers", trust_remote_code=True)
# 3️⃣ Tokenisierung der Texte (hier wird die Spalte "text" genutzt; ggf. anpassen, falls andere Spalten vorhanden sind)
def tokenize_function(examples):
return tokenizer(examples["text"], padding="max_length", truncation=True)
tokenized_datasets = dataset.map(tokenize_function, batched=True)
# 4️⃣ Trainingsparameter setzen
training_args = TrainingArguments(
output_dir="./results",
evaluation_strategy="epoch",
save_strategy="epoch",
per_device_train_batch_size=8,
per_device_eval_batch_size=8,
num_train_epochs=3,
weight_decay=0.01,
logging_dir="./logs",
)
# 5️⃣ Training starten
trainer = Trainer(
model=model,
args=training_args,
train_dataset=tokenized_datasets["train"],
eval_dataset=tokenized_datasets["validation"],
)
trainer.train()
# 6️⃣ Speichern des Modells nach dem Training
model.save_pretrained("./trained_model")
tokenizer.save_pretrained("./trained_model")
print("✅ Training abgeschlossen! Modell gespeichert.")