import torch
from transformers import AutoTokenizer, AutoModelForSequenceClassification

# 加载模型和分词器
model_name = "hfl/chinese-llama-alpaca"
tokenizer = AutoTokenizer.from_pretrained(model_name)
model = AutoModelForSequenceClassification.from_pretrained(model_name, num_labels=2) # 二分类任务

# 加载数据
train_texts = ["这是正样本", "这也是正样本", "这是负样本", "这也是负样本"]
train_labels = [1, 1, 0, 0]

# 分词
train_encodings = tokenizer(train_texts, truncation=True, padding=True)

# 构建Dataset对象
class MyDataset(torch.utils.data.Dataset):
    def __init__(self, encodings, labels):
        self.encodings = encodings
        self.labels = labels

    def __getitem__(self, idx):
        item = {key: torch.tensor(val[idx]) for key, val in self.encodings.items()}
        item['labels'] = torch.tensor(self.labels[idx])
        return item

    def __len__(self):
        return len(self.labels)

train_dataset = MyDataset(train_encodings, train_labels)

# 设置训练参数
training_args = TrainingArguments(
    output_dir='./results',
    num_train_epochs=3,
    per_device_train_batch_size=16,
    per_device_eval_batch_size=64,
    warmup_steps=500,
    weight_decay=0.01,
    logging_dir='./logs',
    logging_steps=10,
    evaluation_strategy='steps',
    eval_steps=50,
    load_best_model_at_end=True
)

# 训练模型
trainer = Trainer(
    model=model,
    args=training_args,
    train_dataset=train_dataset,
    data_collator=default_data_collator,
    compute_metrics=compute_metrics
)

trainer.train()