File size: 6,791 Bytes
b53c3eb |
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 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 |
import torch
import nltk
import numpy as np
import os
import kagglehub
from transformers import AutoTokenizer, AutoModelForCausalLM, AutoModelForSequenceClassification
from bert_score import score as bert_score_calculator
try:
nltk.data.find('tokenizers/punkt')
except nltk.downloader.DownloadError:
nltk.download('punkt')
class LLM_Generator:
def __init__(self, model_handle, device='cuda'):
self.device = device
print(f"Downloading model from Kaggle Hub: {model_handle}")
model_path = kagglehub.model_download(model_handle)
print(f"Model downloaded to: {model_path}")
self.tokenizer = AutoTokenizer.from_pretrained(model_path)
self.model = AutoModelForCausalLM.from_pretrained(
model_path,
torch_dtype="auto",
device_map="auto"
)
def generate(self, prompt, num_samples=1, temperature=0.7, max_new_tokens=150):
messages = [
{"role": "system", "content": "you are a helpful assistant."},
{"role": "user", "content": prompt}
]
text = self.tokenizer.apply_chat_template(
messages,
tokenize=False,
add_generation_prompt=True,
enable_thinking=True
)
model_inputs = self.tokenizer([text] * num_samples, return_tensors="pt").to(self.device)
generated_ids_batch = self.model.generate(
**model_inputs,
max_new_tokens=max_new_tokens,
do_sample=True,
temperature=temperature,
num_return_sequences=num_samples
)
input_ids_len = model_inputs.input_ids.shape[1]
final_responses = []
for generated_ids in generated_ids_batch:
output_ids = generated_ids[input_ids_len:].tolist()
try:
# Find the start of the final content after the "thinking" part
# The token ID 151668 corresponds to the end of the thinking block for Qwen-3
index = len(output_ids) - output_ids[::-1].index(151668)
except ValueError:
index = 0
content = self.tokenizer.decode(output_ids[index:], skip_special_tokens=True).strip("\n")
final_responses.append(content)
return final_responses
class SelfCheckGPT:
def __init__(self, device=None):
if device:
self.device = device
else:
self.device = 'cuda' if torch.cuda.is_available() else 'cpu'
self.nli_tokenizer = None
self.nli_model = None
def _load_nli_model(self):
if self.nli_model is None:
nli_model_name = "microsoft/deberta-v3-large-mnli"
try:
self.nli_tokenizer = AutoTokenizer.from_pretrained(nli_model_name)
self.nli_model = AutoModelForSequenceClassification.from_pretrained(nli_model_name).to(self.device)
except Exception as e:
print(f"Error loading NLI model: {e}")
raise
def _check_bertscore(self, sentences, sample_responses):
all_scores = []
for sent in sentences:
refs = [sent] * len(sample_responses)
cands = sample_responses
_, _, F1 = bert_score_calculator(
cands, refs, lang="en", verbose=False, idf=False, device=self.device
)
avg_bert_score = F1.mean().item()
score = 1.0 - avg_bert_score
all_scores.append(score)
return all_scores
def _check_nli(self, sentences, sample_responses):
self._load_nli_model()
all_scores = []
for sent in sentences:
contradiction_probs = []
for sample in sample_responses:
tokenized_input = self.nli_tokenizer(
sample, sent, return_tensors="pt", truncation=True, max_length=512
).to(self.device)
with torch.no_grad():
logits = self.nli_model(**tokenized_input).logits
entailment_logit = logits[0, self.nli_model.config.label2id['entailment']]
contradiction_logit = logits[0, self.nli_model.config.label2id['contradiction']]
prob_contradiction = torch.exp(contradiction_logit) / (torch.exp(entailment_logit) + torch.exp(contradiction_logit))
contradiction_probs.append(prob_contradiction.item())
avg_contradiction_prob = np.mean(contradiction_probs)
all_scores.append(avg_contradiction_prob)
return all_scores
def check(self, main_response, sample_responses, method='nli'):
sentences = nltk.sent_tokenize(main_response)
if not sentences:
return []
if method.lower() == 'bertscore':
scores = self._check_bertscore(sentences, sample_responses)
elif method.lower() == 'nli':
scores = self._check_nli(sentences, sample_responses)
else:
raise ValueError(f"Invalid method '{method}'. Choose from 'bertscore', 'nli'.")
results = [{"sentence": sent, "score": score} for sent, score in zip(sentences, scores)]
return results
def main():
model_handle = "qwen-lm/qwen-3/transformers/0.6b"
print("Initializing LLM Generator...")
generator = LLM_Generator(model_handle=model_handle)
prompt = "Write a short biography of Neil Armstrong, the first man on the moon. Include the name of the spacecraft he used."
print(f"Generating responses for prompt: '{prompt}'")
responses = generator.generate(prompt, num_samples=6, temperature=0.8, max_new_tokens=150)
main_response = responses[0]
sample_responses = responses[1:]
print("\n--- Generated Main Response ---")
print(main_response)
print("\n--- Generated Sample Responses ---")
for i, r in enumerate(sample_responses):
print(f"{i+1}. {r[:100]}...")
checker = SelfCheckGPT()
print("\n\n--- Running SelfCheckGPT with 'nli' method ---")
nli_results = checker.check(main_response, sample_responses, method='nli')
print("Higher scores suggest a higher probability of being a hallucination.")
for result in nli_results:
print(f"Score: {result['score']:.4f}\tSentence: {result['sentence']}")
print("\n--- Running SelfCheckGPT with 'bertscore' method ---")
bertscore_results = checker.check(main_response, sample_responses, method='bertscore')
print("Higher scores suggest a higher probability of being a hallucination.")
for result in bertscore_results:
print(f"Score: {result['score']:.4f}\tSentence: {result['sentence']}") |