|
|
--- |
|
|
license: cc-by-nc-3.0 |
|
|
base_model: |
|
|
- clue/albert_chinese_tiny |
|
|
--- |
|
|
|
|
|
这是一个用于识别小说章节名称的超轻量级模型，被用于从小说文本中提取章节名称。基于 albert_chinese_tiny。 |
|
|
|
|
|
仓库： [https://github.com/kingzcheung/chapterlm](https://github.com/kingzcheung/chapterlm) |
|
|
|
|
|
```python |
|
|
from transformers import AlbertForSequenceClassification, AutoTokenizer |
|
|
import torch |
|
|
|
|
|
# 加载模型和分词器 |
|
|
model_name = "rkingzhong/chapterlm" # 中文ALBERT-Tiny（仅18MB） |
|
|
tokenizer = AutoTokenizer.from_pretrained(model_name) |
|
|
model = AlbertForSequenceClassification.from_pretrained(model_name, num_labels=2) |
|
|
|
|
|
# print(model) |
|
|
def predict(text): |
|
|
inputs = tokenizer(text, return_tensors="pt", padding=True, truncation=True, max_length=64) |
|
|
with torch.no_grad(): |
|
|
outputs = model(**inputs) |
|
|
print(outputs) |
|
|
return torch.argmax(outputs.logits).item() |
|
|
|
|
|
text = "1、消失的他" |
|
|
|
|
|
pred = predict(text) |
|
|
``` |
