utils.py

import os
import numpy as np
from collections import Counter
from transformers import AutoModelForCausalLM, AutoTokenizer, AutoProcessor
import torch
from peft import PeftModel
import nltk
nltk.download('punkt_tab', quiet=True)
from nltk.tokenize import sent_tokenize, word_tokenize

# Load GPT-2 model and tokenizer for perplexity calculation
_device = "cpu"
_tok = AutoTokenizer.from_pretrained("gpt2")
_model = AutoModelForCausalLM.from_pretrained("gpt2").to(_device)
_model.eval()
if _tok.pad_token is None:
    _tok.pad_token = _tok.eos_token

# Load Gemma model for explanation generation
_gemma_tokenizer = None
_gemma_model = None

HF_TOKEN = os.environ["HF_TOKEN"]

# Set this to where your Unsloth LoRA is stored
BASE_MODEL_REPO = "google/gemma-3-1b-it"
LORA_REPO = "annewaz/gemma-3-1b-it-unsloth-LoRA16"  

def _load_gemma_model():

    """Lazy load Gemma base from bucket + Unsloth LoRA."""
    global _gemma_tokenizer, _gemma_model
    
    if _gemma_tokenizer is None or _gemma_model is None:
        
        # 1. Load processor from your Xet bucket (small files, fast)
        _gemma_tokenizer = AutoTokenizer.from_pretrained(
            BASE_MODEL_REPO,
            token=HF_TOKEN,
            trust_remote_code=True
        )
        
        # 2. Load base model from your bucket
        # Xet streams weights on-demand. device_map="auto" loads directly to GPU/CPU.
        _gemma_model = AutoModelForCausalLM.from_pretrained(
            BASE_MODEL_REPO,
            token = HF_TOKEN,
            device_map="cpu",
            torch_dtype=torch.float32,
            trust_remote_code=True,
        )
        
        # 3. Apply your Unsloth LoRA adapter
        if LORA_REPO:
            print(f"Loading LoRA adapter from {LORA_REPO}...")
            _gemma_model = PeftModel.from_pretrained(_gemma_model, LORA_REPO)
            
            # Optional: merge LoRA into base for faster inference
            # (uses more RAM temporarily during merge, then you can unload base)
            # _gemma_model = _gemma_model.merge_and_unload()
        
        _gemma_model.eval()

   
    
    return _gemma_tokenizer, _gemma_model

# def _load_gemma_model():
#     """Lazy load Gemma model for explanation generation."""
#     global _gemma_processor, _gemma_model
#     if _gemma_processor is None or _gemma_model is None:
#         _gemma_processor = AutoProcessor.from_pretrained("google/gemma-4-E4B-it")
#         _gemma_model = AutoModelForMultimodalLM.from_pretrained("google/gemma-4-E4B-it")
#         _gemma_model.eval()
#     return _gemma_processor, _gemma_model


def compute_burstiness(text):
    """
    Compute burstiness score for a single text.
    Burstiness = variance/mean - 1 of word frequencies.
    Higher values indicate more bursty (human-like) writing.
    """
    words = [w.lower() for w in word_tokenize(str(text))]
    if len(words) == 0:
        return 0.0
    
    freqs = list(Counter(words).values())
    mu = np.mean(freqs)
    if mu == 0:
        return 0.0
    
    return float(np.var(freqs) / mu - 1.0)


def compute_ttr(text):
    """
    Compute Type-Token Ratio for a single text.
    TTR = unique words / total words.
    Higher values indicate richer vocabulary.
    """
    words = [w.lower() for w in word_tokenize(str(text))]
    if len(words) == 0:
        return 0.0
    
    return float(len(set(words)) / len(words))


def compute_cv_sentence_length(text):
    """
    Compute Coefficient of Variation of sentence lengths.
    CV = std(sentence_lengths) / mean(sentence_lengths).
    Higher values indicate more variation in sentence structure.
    """
    sents = sent_tokenize(str(text))
    lens = [len(word_tokenize(x)) for x in sents]
    
    if len(lens) <= 1:
        return 0.0
    
    mu = np.mean(lens)
    if mu == 0:
        return 0.0
    
    return float(np.std(lens) / mu)


def compute_perplexity(text, max_length=64):
    """
    Compute perplexity of text using GPT-2.
    Lower perplexity indicates more predictable (likely AI-generated) text.
    """
    text = str(text)
    
    with torch.no_grad():
        enc = _tok(text, padding=True, truncation=True,
                   max_length=max_length, return_tensors="pt").to(_device)
        input_ids = enc["input_ids"]
        attention_mask = enc["attention_mask"]
        labels = input_ids.clone()
        labels[attention_mask == 0] = -100
        
        outputs = _model(input_ids=input_ids, attention_mask=attention_mask, labels=labels)
        logits = outputs.logits
        
        shift_logits = logits[..., :-1, :].contiguous()
        shift_labels = labels[..., 1:].contiguous()
        
        loss_fct = torch.nn.CrossEntropyLoss(reduction='none', ignore_index=-100)
        per_token_loss = loss_fct(shift_logits.view(-1, shift_logits.size(-1)), 
                                  shift_labels.view(-1))
        per_token_loss = per_token_loss.view(shift_labels.shape)
        
        seq_mask = (shift_labels != -100).float()
        per_sample_loss = (per_token_loss * seq_mask).sum(dim=1) / seq_mask.sum(dim=1).clamp(min=1)
        perplexity = torch.exp(per_sample_loss).cpu().numpy().flatten()[0]
    
    return float(perplexity)


def compute_all_linguistic_features(text):
    """
    Compute all four linguistic features at once.
    Returns a dictionary with burstiness, TTR, CV_sentence_len, and perplexity.
    """
    return {
        'burstiness': compute_burstiness(text),
        'TTR': compute_ttr(text),
        'CV_sentence_len': compute_cv_sentence_length(text),
        'perplexity': compute_perplexity(text)
    }


def text_input_generate(text, prediction, confidence, linguistic_features):
    """
    Generate a combined text input for later use.
    
    Args:
        text: The original input text
        prediction: The prediction label (e.g., "AI-generated" or "Human-written")
        confidence: The confidence score (e.g., 0.8542)
        linguistic_features: Dictionary with burstiness, TTR, CV_sentence_len, and perplexity
    
    Returns:
        A combined text string with all information
    """
    combined_text = f"""Input Text:
{text}

Prediction: {prediction}
Confidence: {confidence:.4f} ({confidence:.2%})

Linguistic Features:
- Burstiness: {linguistic_features['burstiness']:.4f}
- TTR (Type-Token Ratio): {linguistic_features['TTR']:.4f}
- CV (Coefficient of Variation of Sentence Length): {linguistic_features['CV_sentence_len']:.4f}
- Perplexity: {linguistic_features['perplexity']:.4f}
"""
    return combined_text


def generate_explanation_with_gemma(combined_text):
    """
    
    
    Args:
        combined_text: The combined text with input, prediction, and features
    
    Returns:
        Generated explanation text
    """
    try:
        tokenizer, model = _load_gemma_model()
        
        # Prompt
        messages = [
            {"role": "system", "content": "Please explain why the text is either AI generated or human written within 100 words at most?"},
            {"role": "user", "content": combined_text},
        ]
        
        # Process input
        text = tokenizer.apply_chat_template(
            messages, 
            tokenize=False, 
            add_generation_prompt=True, 
            # enable_thinking=False
        )
        inputs = tokenizer(text=text, return_tensors="pt")
        inputs = {k: v.to(model.device) for k, v in inputs.items()}
        input_len = inputs["input_ids"].shape[-1]
        
        # Generate output
        outputs = model.generate(**inputs, max_new_tokens=256)
        response = tokenizer.decode(outputs[0][input_len:], skip_special_tokens=True)
        
        # Parse output
        # explanation = tokenizer.parse_response(response)
        
        return response
    except Exception as e:
        return f"Error generating explanation: {str(e)}"