Create main_kaggle.py

main_kaggle.py

@@ -0,0 +1,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']}")