import json
import torch
from huggingface_hub import hf_hub_download
from peft import AutoPeftModelForSequenceClassification
from transformers import AutoTokenizer
model_id = "chris0809/memoperator-0.6b-memory-write-gate"
tokenizer = AutoTokenizer.from_pretrained(model_id)
if tokenizer.pad_token_id is None:
tokenizer.pad_token = tokenizer.eos_token
model = AutoPeftModelForSequenceClassification.from_pretrained(model_id)
model.config.pad_token_id = tokenizer.pad_token_id
metadata = json.loads(open(hf_hub_download(model_id, "classifier_metadata.json"), encoding="utf-8").read())
inputs = tokenizer("以后给我写周报时先写结论。", return_tensors="pt")
with torch.inference_mode():
save_probability = torch.softmax(model(**inputs).logits, dim=-1)[0, 1].item()
decision = "SAVE" if save_probability >= metadata["threshold"] else "SKIP"