-
Notifications
You must be signed in to change notification settings - Fork 9
Expand file tree
/
Copy pathdemo_v2.py
More file actions
53 lines (44 loc) · 2.15 KB
/
demo_v2.py
File metadata and controls
53 lines (44 loc) · 2.15 KB
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
import torch
from torch.nn.functional import sigmoid
from transformers import AutoTokenizer, AutoModelForSequenceClassification
from peft import AutoPeftModelForSequenceClassification
class NAIPv2:
def __init__(self, model_path: str, device: str = 'cuda', max_length: int = 512):
self.device = torch.device(device if torch.cuda.is_available() else 'cpu')
self.max_length = max_length
# Load model
self.model = AutoPeftModelForSequenceClassification.from_pretrained(
model_path,
load_in_8bit=True,
device_map={"": torch.cuda.current_device()} if self.device.type == "cuda" else "auto",
num_labels=1
).eval()
# Load tokenizer and set pad_token
self.tokenizer = AutoTokenizer.from_pretrained(model_path)
self.tokenizer.pad_token = self.tokenizer.eos_token
self.model.config.pad_token_id = self.tokenizer.pad_token_id
# ⚠️ DO NOT CHANGE PROMPT
self.prompt_template = (
"Given a research paper, Title: {title}\nAbstract: {abstract}\nEvaluate the quality of this paper:"
)
def score(self, title: str, abstract: str) -> float:
prompt = self.prompt_template.format(title=title.strip(), abstract=abstract.strip())
inputs = self.tokenizer(prompt, return_tensors='pt', padding=True, truncation=True,
max_length=self.max_length).to(self.device)
with torch.no_grad():
logits = self.model(**inputs).logits # You may scale logits depending on your needs
score = logits.view(-1).item()
return score
if __name__ == "__main__":
model_path = r"path_to_the_v2_adapter_dir"
scorer = NAIPv2(model_path=model_path, device='cuda')
print("🎯 Enter paper title and abstract, press Enter to get the predicted score (Ctrl+C to quit)\n")
while True:
try:
title = input("📄 Title: ")
abstract = input("📑 Abstract: ")
score = scorer.score(title, abstract)
print(f"🔮 Predicted Score: {score:.4f}\n")
except KeyboardInterrupt:
print("\nExiting.")
break