| """用法：python predict.py '商品标题'""" |
|
|
| import argparse |
| import json |
| from pathlib import Path |
|
|
| import torch |
| from transformers import AutoTokenizer |
|
|
| from model import BertTitleClassifier |
|
|
|
|
| def main(): |
| parser = argparse.ArgumentParser(description="商品标题分类") |
| parser.add_argument("text", help="待分类的商品标题") |
| args = parser.parse_args() |
| model_dir = Path(__file__).resolve().parent |
| settings = json.loads((model_dir / "inference_config.json").read_text(encoding="utf-8")) |
| labels = json.loads((model_dir / "labels.json").read_text(encoding="utf-8")) |
| device = "cuda" if torch.cuda.is_available() else "cpu" |
| model = BertTitleClassifier.from_pretrained(model_dir, device=device) |
| tokenizer = AutoTokenizer.from_pretrained(model_dir, local_files_only=True) |
| inputs = tokenizer( |
| args.text, return_tensors="pt", padding="max_length", |
| truncation=True, max_length=settings["max_length"], |
| ) |
| with torch.no_grad(): |
| logits = model(inputs["input_ids"].to(device), inputs["attention_mask"].to(device)) |
| pred_id = logits.argmax(dim=1).item() |
| print(json.dumps({"id": pred_id, "label": labels["id2label"][str(pred_id)]}, ensure_ascii=False)) |
|
|
|
|
| if __name__ == "__main__": |
| main() |
|
|
