Views
No views yet
1import os;
2os.environ["HF_ENDPOINT"] = "https://hf-mirror.com"
3
4import json
5import torch
6from huggingface_hub import hf_hub_download
7from transformers import AutoTokenizer, AutoModel
8
9model_id = "xulab-research/patent-classifier-4B"
10device = "cuda" if torch.cuda.is_available() else "cpu"
11
12tokenizer = AutoTokenizer.from_pretrained(model_id, trust_remote_code=True)
13model = AutoModel.from_pretrained(model_id, trust_remote_code=True).to(device).eval()
14
15# 加载最优阈值
16th_path = hf_hub_download(repo_id=model_id, filename="optimal_thresholds.json")
17with open(th_path, "r", encoding="utf-8") as f:
18 th_dict = json.load(f)
19
20id2label = model.config.id2label
21label2th = {id2label[i]: th_dict[str(i)] for i in range(len(id2label))}
22
23texts = [
24 "这里是一段专利摘要示例文本 A",
25 "这里是另一段专利摘要示例文本 B",
26]
27
28inputs = tokenizer(texts, padding=True, truncation=True, max_length=256, return_tensors="pt").to(device)
29with torch.no_grad():
30 logits = model(**inputs) # 形状: [batch, num_labels]
31
32probs = torch.sigmoid(logits).cpu().tolist()
33
34predicted_labels = []
35for row in probs:
36 labels = []
37 for i, p in enumerate(row):
38 label = id2label[i]
39 if p >= label2th[label]:
40 labels.append(label)
41 predicted_labels.append(labels)
42
43print(predicted_labels)
44# 例如: [["LABEL_0", "LABEL_3"], ["LABEL_2"]]
45 1# parquet文件结构
2df = pd.DataFrame({
3 "摘要文本": ["这是一个专利摘要文本...", "另一个专利摘要..."]
4})