Spaces:
Sleeping
Sleeping
File size: 2,224 Bytes
8ab2445 |
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 |
from transformers import AutoTokenizer, AutoModelForSequenceClassification
import torch
class GroundednessChecker:
def __init__(self, model_path="./grounding_detector"):
self.tokenizer = AutoTokenizer.from_pretrained(model_path)
self.model = AutoModelForSequenceClassification.from_pretrained(model_path)
self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
self.model.to(self.device)
def check(self, question: str, answer: str, context: str) -> dict:
"""Check if answer is grounded in context"""
inputs = self.tokenizer(
question,
answer + " [SEP] " + context,
padding=True,
truncation=True,
max_length=512,
return_tensors="pt"
).to(self.device)
with torch.no_grad():
outputs = self.model(**inputs)
probs = torch.nn.functional.softmax(outputs.logits, dim=-1)
return {
"is_grounded": bool(torch.argmax(probs)),
"confidence": probs[0][1].item(),
"details": {
"question": question,
"answer": answer,
"context_snippet": context[:200] + "..." if len(context) > 200 else context
}
}
# Usage Example
if __name__ == "__main__":
# Initialize checker
checker = GroundednessChecker()
# Example from banking PDS
context = """
Premium Savings Account Terms:
- Annual Percentage Yield (APY): 4.25%
- Minimum opening deposit: $1,000
- Monthly maintenance fee: $5 (waived if daily balance >= $1,000)
- Maximum withdrawals: 6 per month
"""
# Grounded example
grounded_result = checker.check(
question="What is the minimum opening deposit?",
answer="$1,000",
context=context
)
print("Grounded Result:", grounded_result)
# Ungrounded example
ungrounded_result = checker.check(
question="What is the monthly maintenance fee?",
answer="$10 monthly charge",
context=context
)
print("Ungrounded Result:", ungrounded_result)
|