Spaces:
Sleeping
Sleeping
File size: 3,347 Bytes
006a1d9 7ff3e94 006a1d9 7ff3e94 006a1d9 7ff3e94 006a1d9 7ff3e94 006a1d9 7ff3e94 006a1d9 7ff3e94 006a1d9 7ff3e94 006a1d9 7ff3e94 41c2d4c 7ff3e94 41c2d4c 006a1d9 41c2d4c 006a1d9 41c2d4c | 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 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 | """Inference helpers for Dialectica."""
ORDERED_LABELS = ["Surface", "Mechanistic", "Critical"]
class CognitiveClassifier:
"""DistilBERT question classifier."""
def __init__(self, product_config):
"""Load model and tokenizer."""
import torch
from transformers import (
AutoModelForSequenceClassification,
AutoTokenizer,
)
self.cfg = product_config
self.torch = torch
self.tokenizer = AutoTokenizer.from_pretrained(self.cfg.classifier_dir)
self.model = AutoModelForSequenceClassification.from_pretrained(
self.cfg.classifier_dir
)
self.model.eval()
self.id2label = self.model.config.id2label or {
i: label for i, label in enumerate(ORDERED_LABELS)
}
def classify(self, question):
"""Predict level and confidence."""
inputs = self.tokenizer(
question,
truncation=True,
max_length=self.cfg.max_question_length,
return_tensors="pt",
)
with self.torch.no_grad():
logits = self.model(**inputs).logits
probs = self.torch.softmax(logits, dim=1)[0]
predicted_id = int(self.torch.argmax(probs))
return {
"level": self.id2label[predicted_id],
"confidence": float(probs[predicted_id]),
}
class ConceptMatcher:
"""Embedding-based concept matcher."""
def __init__(self, product_config):
"""Load embedding model."""
from sentence_transformers import SentenceTransformer
self.cfg = product_config
self.model = SentenceTransformer(self.cfg.embed_model)
self.concepts = []
self.concept_embeddings = None
def set_concepts(self, concepts):
"""Cache concept embeddings."""
self.concepts = concepts
if concepts:
self.concept_embeddings = self.model.encode(
concepts, convert_to_tensor=True
)
else:
self.concept_embeddings = None
def match(self, question):
"""Return up to two matching concepts."""
from sentence_transformers import util
if not self.concepts or self.concept_embeddings is None:
return []
query = self.model.encode(question, convert_to_tensor=True)
scores = util.cos_sim(query, self.concept_embeddings)[0]
ranked = sorted(
range(len(self.concepts)),
key=lambda i: float(scores[i]),
reverse=True,
)
threshold = self.cfg.concept_match_threshold
# Keep a second concept if it's close to top or strong enough.
runner_up_margin = 0.05
runner_up_absolute = 0.45
top_index = ranked[0]
top_score = float(scores[top_index])
if top_score < threshold:
return []
matched = [self.concepts[top_index]]
if len(ranked) > 1:
second_index = ranked[1]
second_score = float(scores[second_index])
close_to_top = top_score - second_score <= runner_up_margin
strong_alone = second_score >= runner_up_absolute
if second_score >= threshold and (close_to_top or strong_alone):
matched.append(self.concepts[second_index])
return matched
|