Views
No views yet
1from transformers import AlbertForSequenceClassification, AutoTokenizer
2import torch
3
4# 加载模型和分词器
5model_name = "rkingzhong/chapterlm" # 中文ALBERT-Tiny（仅18MB）
6tokenizer = AutoTokenizer.from_pretrained(model_name)
7model = AlbertForSequenceClassification.from_pretrained(model_name, num_labels=2)
8
9# print(model)
10def predict(text):
11 inputs = tokenizer(text, return_tensors="pt", padding=True, truncation=True, max_length=64)
12 with torch.no_grad():
13 outputs = model(**inputs)
14 print(outputs)
15 return torch.argmax(outputs.logits).item()
16
17text = "1、消失的他"
18
19pred = predict(text)