Views
No views yet
1import torch
2from transformers import AutoModel
3from torch import nn
4from huggingface_hub import hf_hub_download
5from transformers import AutoModel, AutoTokenizer
6tokenizer = AutoTokenizer.from_pretrained("CMCenjoyer/deberta-trace")
7
8
9class DebertaTrace(nn.Module):
10 def __init__(self, base_model):
11 super().__init__()
12 self.base = base_model
13 hid = base_model.config.hidden_size
14 self.rel_head = nn.Linear(hid,1)
15 self.util_head = nn.Linear(hid,1)
16 self.adh_head = nn.Linear(hid,1)
17
18 def forward(self, input_ids, attention_mask):
19 out = self.base(input_ids=input_ids, attention_mask=attention_mask)
20 hs = out.last_hidden_state
21 return {
22 'logits_relevance': self.rel_head(hs).squeeze(-1),
23 'logits_utilization': self.util_head(hs).squeeze(-1),
24 'logits_adherence': self.adh_head(hs).squeeze(-1)
25 }
26
27base_model = AutoModel.from_pretrained("CMCenjoyer/deberta-trace")
28model = DebertaTrace(base_model)
29# heads_weights.p в локальный кэш
30file_path = hf_hub_download(repo_id="CMCenjoyer/deberta-trace", filename="heads_weights.pt")
31heads_weights = torch.load(file_path, weights_only=True)
32model.rel_head.load_state_dict(heads_weights['rel_head'])
33model.util_head.load_state_dict(heads_weights['util_head'])
34model.adh_head.load_state_dict(heads_weights['adh_head'])
35def preprocess(example, max_length=512):
36 '''
37 Препроцессим входной элемент в маску контекста, маску ответва и input_ids + attention_mask
38 '''
39 question_ids = tokenizer.encode(example["question"], add_special_tokens=False)
40
41 doc_ids = []
42 for doc in example["documents_sentences"]:
43 for _, sent in doc:
44 tokens = tokenizer.encode(sent, add_special_tokens=False)
45 doc_ids += tokens
46
47 response_ids = tokenizer.encode(example["response"], add_special_tokens=False)
48
49 sep_id = tokenizer.sep_token_id
50 input_ids = question_ids + [sep_id] + doc_ids + [sep_id] + response_ids
51
52 context_mask = [0] * (len(question_ids) + 1) + [1] * len(doc_ids) + [0] + [0] * len(response_ids)
53 response_mask = [0] * (len(question_ids) + len(doc_ids) + 2) + [1] * len(response_ids)
54
55 if len(input_ids) > max_length:
56 input_ids = input_ids[:max_length]
57 context_mask = context_mask[:max_length]
58 response_mask = response_mask[:max_length]
59
60 return {
61 "input_ids": torch.tensor(input_ids, dtype=torch.long),
62 "attention_mask": torch.tensor([1] * len(input_ids), dtype=torch.long),
63 "context_mask": torch.tensor(context_mask, dtype=torch.bool),
64 "response_mask": torch.tensor(response_mask, dtype=torch.bool),
65 }
66def compute_trace_metrics_inference(logits, masks, threshold=0.5):
67 '''
68 подсчет метрик TRACE для каждого элемента батча(все батчи должны быть фиксированной одной длины)
69 '''
70 rel_pred = (torch.sigmoid(logits['logits_relevance'].detach().cpu()) > threshold)
71 util_pred = (torch.sigmoid(logits['logits_utilization'].detach().cpu())> threshold)
72 adh_pred = (torch.sigmoid(logits['logits_adherence'].detach().cpu()) > threshold)
73
74 ctx_m = masks['context_mask'].detach().cpu()
75 resp_m = masks['response_mask'].detach().cpu()
76
77 def rate(pred, mask):
78 # sum(pred & mask) / sum(mask)
79 num = (pred & mask).sum(dim=1).float()
80 den = mask.sum(dim=1).float().clamp(min=1)
81 return num.div(den)
82
83 relevance_rate = rate(rel_pred, ctx_m)
84 utilization_rate = rate(util_pred, ctx_m)
85 adherence_rate = rate(adh_pred, resp_m)
86
87 # completeness: из релевантных предсказаний — сколько ещё и util
88 num_ru = (rel_pred & util_pred & ctx_m).sum(dim=1).float()
89 den_r = rel_pred.sum(dim=1).float().clamp(min=1)
90 completeness = num_ru.div(den_r)
91
92 return {
93 'relevance_rate': relevance_rate,
94 'utilization_rate': utilization_rate,
95 'adherence_rate': adherence_rate,
96 'completeness': completeness
97 }
98from datasets import load_dataset
99ds = load_dataset("rungalileo/ragbench", "delucionqa")
100ex = preprocess(ds['train'][9])
101model.eval()
102with torch.no_grad():
103 outputs = model(ex["input_ids"].unsqueeze(0), ex["attention_mask"].unsqueeze(0))
104 batch_metrics = compute_trace_metrics_inference(outputs, {'context_mask': ex["context_mask"].unsqueeze(0) , 'response_mask':ex["response_mask"].unsqueeze(0)})
105batch_metrics