在低配置电脑上使用TinyBERT训练并部署产品编号提取模型

pip install torch transformers datasets seqeval scikit-learn

订单号是：B1234567，它是产品A123的编号。

订单号是：O B-PRODUCT I-PRODUCT I-PRODUCT O，它是产品O A-PRODUCT I-PRODUCT 的编号。

[
  {
    "tokens": ["订单号", "是", "B1234567", "它", "是", "产品A123", "的", "编号", "。"],
    "labels": ["O", "O", "B-PRODUCT", "O", "O", "B-PRODUCT", "O", "O", "O"]
  },
  {
    "tokens": ["产品", "编号", "是", "XyZ-001", "，", "请", "记得", "标注", "。"],
    "labels": ["O", "O", "O", "B-PRODUCT", "O", "O", "O", "O", "O"]
  }
]

from transformers import BertTokenizer, BertForTokenClassification
from transformers import Trainer, TrainingArguments

# 加载预训练模型和Tokenizer
tokenizer = BertTokenizer.from_pretrained('huawei-noah/TinyBERT_General_4L_312D')
model = BertForTokenClassification.from_pretrained('huawei-noah/TinyBERT_General_4L_312D', num_labels=3)

# 设置标签
label_list = ["O", "B-PRODUCT", "I-PRODUCT"]
label2id = {label: i for i, label in enumerate(label_list)}
id2label = {i: label for i, label in enumerate(label_list)}

model.config.label2id = label2id
model.config.id2label = id2label

from datasets import load_dataset

# 加载并处理数据
train_data = load_dataset('json', data_files='train_data.json')['train']

def tokenize_and_align_labels(examples):
    tokenized_inputs = tokenizer(examples['tokens'], truncation=True, padding='max_length', is_split_into_words=True)
    labels = examples['labels']
    label_ids = [label2id[label] for label in labels]
    tokenized_inputs["labels"] = label_ids
    return tokenized_inputs

train_data = train_data.map(tokenize_and_align_labels, batched=True)

training_args = TrainingArguments(
    output_dir='./results',
    num_train_epochs=3,
    per_device_train_batch_size=8,
    per_device_eval_batch_size=8,
    logging_dir='./logs',
    evaluation_strategy="epoch"
)

trainer = Trainer(
    model=model,
    args=training_args,
    train_dataset=train_data,
)

trainer.train()

test_data = load_dataset('json', data_files='test_data.json')['test']
test_data = test_data.map(tokenize_and_align_labels, batched=True)

# 使用Trainer进行评估
trainer.evaluate(test_data)

model.save_pretrained('./product_number_model')
tokenizer.save_pretrained('./product_number_model')

from fastapi import FastAPI
from pydantic import BaseModel

app = FastAPI()

# 加载模型
model = BertForTokenClassification.from_pretrained('./product_number_model')
tokenizer = BertTokenizer.from_pretrained('./product_number_model')

class TextInput(BaseModel):
    text: str

@app.post("/extract_product_number")
def extract_product_number(input: TextInput):
    tokens = tokenizer(input.text, return_tensors="pt")
    outputs = model(**tokens)
    predictions = outputs.logits.argmax(dim=-1).squeeze().tolist()

    # 映射预测标签
    labels = [model.config.id2label[pred] for pred in predictions]
    return {"text": input.text, "labels": labels}

uvicorn main:app --reload

{
  "text": "订单号是 B1234567，它是产品A123的编号"
}