😽

自分用trl学習手順 Copilot生成

2024/04/30に公開
  1. モデルのロード:
from transformers import AutoModelForCausalLM, AutoTokenizer
model = AutoModelForCausalLM.from_pretrained("モデル名")
tokenizer = AutoTokenizer.from_pretrained("モデル名")
  1. データセットのロードとトークナイズ:
from datasets import load_dataset
dataset = load_dataset("データセット名")
dataset = dataset.map(lambda examples: tokenizer(examples["text"], truncation=True, padding='max_length'), batched=True)
  1. オプティマイザとメトリクスの設定:
from transformers import AdamW
from sklearn.metrics import accuracy_score

optimizer = AdamW(model.parameters(), lr=1e-5)

def compute_metrics(pred):
    labels = pred.label_ids
    preds = pred.predictions.argmax(-1)
    acc = accuracy_score(labels, preds)
    return {
        'accuracy': acc,
    }
  1. DPOTrainerの設定と学習:
from trl.trainer import DPOTrainer
trainer = DPOTrainer(
    model=model,
    data_collator=DataCollatorForCompletionOnlyLM(tokenizer=tokenizer, prompt_token_id=tokenizer.bos_token_id, end_token_id=tokenizer.eos_token_id),
    train_dataset=dataset["train"],
    eval_dataset=dataset["validation"],
    compute_metrics=compute_metrics,
    optimizers=(optimizer, None),
)
trainer.train()
  1. 推論:
prompt = "あなたの好きな食べ物は何ですか?"
inputs = tokenizer.encode(prompt, return_tensors='pt')
outputs = model.generate(inputs, max_length=100, temperature=0.7)
print(tokenizer.decode(outputs[0]))

このコードは、訓練したモデルを使用して、指定したプロンプトに対する推論を行います。

  1. モデルの保存:
model.save_pretrained("保存するパス")
tokenizer.save_pretrained("保存するパス")

このコードは、訓練したモデルとトークナイザを指定したパスに保存します。

Discussion