微调ALBERT问题并用HuggingFace回答

from transformers import AlbertTokenizer, AlbertForQuestionAnswering, Trainer, TrainingArguments
import torch

# 加载预训练模型和分词器
model_name = 'albert-base-v2'
tokenizer = AlbertTokenizer.from_pretrained(model_name)
model = AlbertForQuestionAnswering.from_pretrained(model_name)

# 准备数据集
class QADataset(torch.utils.data.Dataset):
    def __init__(self, encodings):
        self.encodings = encodings

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

    def __len__(self):
        return len(self.encodings['input_ids'])

# 示例数据
data = {
    'question': 'What is the capital of France?',
    'context': 'The capital of France is Paris.',
    'answers': [{'answer_start': 19, 'text': 'Paris'}]
}

# 编码数据
inputs = tokenizer(data['question'], data['context'], truncation=True, padding='max_length', max_length=512)
encodings = {key: inputs[key].squeeze() for key in inputs}
dataset = QADataset(encodings)

# 设置训练参数
training_args = TrainingArguments(
    output_dir='./results',
    num_train_epochs=3,
    per_device_train_batch_size=8,
    per_device_eval_batch_size=8,
    warmup_steps=500,
    weight_decay=0.01,
    logging_dir='./logs',
)

# 创建Trainer
trainer = Trainer(
    model=model,
    args=training_args,
    train_dataset=dataset,
    eval_dataset=dataset
)

# 训练模型
trainer.train()

# 进行问答
def answer_question(question, context):
    inputs = tokenizer(question, context, truncation=True, padding='max_length', max_length=512, return_tensors='pt')
    outputs = model(**inputs)
    answer_start = torch.argmax(outputs.start_logits)
    answer_end = torch.argmax(outputs.end_logits) + 1
    answer = tokenizer.convert_tokens_to_string(tokenizer.convert_ids_to_tokens(inputs['input_ids'][0][answer_start:answer_end]))
    return answer

# 测试问答
question = 'What is the capital of France?'
context = 'The capital of France is Paris.'
print(answer_question(question, context))