rkingzhong/chapterlm

This is a super lightweight model for identifying novel chapter names and is used to extract chapter names from novel texts.

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)