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)}"