README.md

metadata

license: cc-by-nc-3.0
base_model:
  - clue/albert_chinese_tiny

这是一个用于识别小说章节名称的超轻量级模型，被用于从小说文本中提取章节名称。基于 albert_chinese_tiny。

仓库： https://github.com/kingzcheung/chapterlm

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)