Upload scientific_model_inference2.py with huggingface_hub

scientific_model_inference2.py

@@ -0,0 +1,989 @@
+#!/usr/bin/env python3
+"""
+Scientific Summarization Model Inference Module - FIXED VERSION
+Fixed generation errors and improved title quality
+"""
+
+import torch
+import torch.nn as nn
+import pandas as pd
+import numpy as np
+import pickle
+import json
+import re
+from pathlib import Path
+from sentence_transformers import SentenceTransformer
+from transformers import AutoTokenizer, AutoModelForCausalLM
+from peft import get_peft_model, LoraConfig, TaskType
+from typing import Dict, List, Tuple, Optional
+from datetime import datetime
+import csv
+from collections import defaultdict, Counter
+from tqdm import tqdm
+import unicodedata
+import hashlib
+import os
+import gc
+import warnings
+
+# Suppress transformer warnings
+warnings.filterwarnings("ignore", category=UserWarning, module="transformers")
+
+# SPEED OPTIMIZATION: Enhanced environment setup for RTX 3080
+os.environ["CUDA_VISIBLE_DEVICES"] = "0"
+os.environ["NCCL_P2P_DISABLE"] = "0"
+os.environ["NCCL_IB_DISABLE"] = "0"
+os.environ["ACCELERATE_DEVICE_PLACEMENT"] = "false"
+os.environ["PYTORCH_CUDA_ALLOC_CONF"] = "max_split_size_mb:512,expandable_segments:True"
+
+# SPEED OPTIMIZATION: Enable all performance optimizations
+torch.backends.cuda.matmul.allow_tf32 = True
+torch.backends.cudnn.allow_tf32 = True
+torch.backends.cudnn.benchmark = True
+torch.backends.cudnn.deterministic = False
+torch._dynamo.config.suppress_errors = True
+
+class Sbert2Prompt(nn.Module):
+    """Prompt generator from SBERT embeddings - matching training architecture"""
+    def __init__(self, sbert_dim, llama_hidden_dim, prompt_length=24):  # Using 24 from training
+        super().__init__()
+        self.prompt_length = prompt_length
+        self.llama_hidden_dim = llama_hidden_dim
+
+        self.projection = nn.Sequential(
+            nn.Linear(sbert_dim, llama_hidden_dim * 2),
+            nn.GELU(),
+            nn.Dropout(0.1),
+            nn.Linear(llama_hidden_dim * 2, llama_hidden_dim * prompt_length)
+        )
+
+    def forward(self, sbert_emb):
+        B = sbert_emb.size(0)
+        out = self.projection(sbert_emb)
+        return out.view(B, self.prompt_length, self.llama_hidden_dim)
+
+def normalize_characters(text):
+    """Normalize various Unicode characters to standard ASCII equivalents"""
+    if not isinstance(text, str):
+        return str(text)
+
+    # Normalize space characters
+    space_chars = ['\xa0', '\u2000', '\u2001', '\u2002', '\u2003', '\u2004', '\u2005', '\u2006', '\u2007', '\u2008', '\u2009', '\u200a', '\u202f', '\u205f', '\u3000']
+    for space in space_chars:
+        text = text.replace(space, ' ')
+
+    # Normalize single quotes
+    single_quotes = [''', ''', '‛', '′', '‹', '›', '‚', '‟']
+    for quote in single_quotes:
+        text = text.replace(quote, "'")
+
+    # Normalize double quotes
+    double_quotes = ['"', '"', '„', '‟', '«', '»', '〝', '〞', '〟', '＂']
+    for quote in double_quotes:
+        text = text.replace(quote, '"')
+
+    # Remove or normalize any remaining special characters
+    text = unicodedata.normalize('NFKD', text)
+    return text
+
+def clean_text(text):
+    """Clean and validate text data"""
+    if not text or str(text) in ['nan', 'None', '']:
+        return ""
+
+    text = normalize_characters(str(text))
+    text = re.sub(r'\s+', ' ', text).strip()
+    return text
+
+class ScientificModelInference:
+    """Main inference class with fixed generation and better titles"""
+
+    def __init__(self, model_dir: str, device: str = "auto"):
+        """
+        Initialize the inference model with enhanced generation capabilities
+
+        Args:
+            model_dir: Path to saved model directory
+            device: Device to use ('auto', 'cuda', 'cpu')
+        """
+        self.model_dir = Path(model_dir)
+        self.device = device if device != "auto" else ("cuda" if torch.cuda.is_available() else "cpu")
+
+        # Load configuration
+        with open(self.model_dir / "config.json", 'r') as f:
+            self.config = json.load(f)
+
+        # ENHANCED: Update prompt length to match training (24)
+        if 'prompt_length' in self.config:
+            self.config['prompt_length'] = 24  # Match training configuration
+
+        print(f"🔧 Loading model on device: {self.device}")
+        self._load_models()
+
+        # Store keywords for title generation context
+        self._last_keywords = []
+        self._last_abstracts = []  # ENHANCED: Store abstracts for better context
+
+        # ENHANCED: Track title generation patterns and word frequency to avoid repetition
+        self._title_patterns_used = Counter()
+        self._title_word_frequency = Counter()  # Track word usage across all titles
+
+        # SPEED OPTIMIZATION: Compile model for faster inference if supported
+        self._optimize_models()
+
+    def _load_models(self):
+        """Load all required models with speed optimizations"""
+        # SPEED OPTIMIZATION: Load SBERT model with optimizations
+        print("📊 Loading SBERT model with optimizations...")
+        self.sbert_model = SentenceTransformer(self.config['sbert_model_name'])
+        self.sbert_model = self.sbert_model.to(self.device)
+        self.sbert_model.eval()
+
+        # SPEED OPTIMIZATION: Disable gradients for SBERT
+        for param in self.sbert_model.parameters():
+            param.requires_grad = False
+
+        # Load tokenizer with optimizations
+        print("🔤 Loading tokenizer...")
+        self.tokenizer = AutoTokenizer.from_pretrained(self.model_dir / "model")
+        if self.tokenizer.pad_token is None:
+            self.tokenizer.pad_token = self.tokenizer.eos_token
+
+        # SPEED OPTIMIZATION: Load main model with better memory settings
+        print("🧠 Loading language model with enhanced generation support...")
+        self.model = AutoModelForCausalLM.from_pretrained(
+            self.model_dir / "model",
+            torch_dtype=torch.float16,
+            device_map="auto" if self.device == "cuda" else None,
+            low_cpu_mem_usage=True,
+            use_cache=True,
+            attn_implementation="flash_attention_2" if hasattr(torch.nn, 'scaled_dot_product_attention') else "eager"
+        )
+        self.model.eval()
+
+        # SPEED OPTIMIZATION: Disable gradients for inference
+        for param in self.model.parameters():
+            param.requires_grad = False
+
+        # Load prompt generator with correct architecture
+        print("⚡ Loading prompt generator (24 tokens)...")
+        self.prompt_generator = Sbert2Prompt(
+            self.config['embedding_dim'],
+            self.config['llama_hidden_dim'],
+            24  # Match training prompt length
+        )
+        self.prompt_generator.load_state_dict(
+            torch.load(self.model_dir / "prompt_generator.pt", map_location=self.device, weights_only=False)
+        )
+        self.prompt_generator = self.prompt_generator.to(self.device, dtype=torch.float16)
+        self.prompt_generator.eval()
+
+        # SPEED OPTIMIZATION: Disable gradients for prompt generator
+        for param in self.prompt_generator.parameters():
+            param.requires_grad = False
+
+        print("✅ All models loaded with enhanced generation support!")
+
+    def _optimize_models(self):
+        """Apply additional speed optimizations"""
+        try:
+            # SPEED OPTIMIZATION: Try to compile models for faster inference (PyTorch 2.0+)
+            if hasattr(torch, 'compile') and torch.cuda.is_available():
+                print("🚀 Applying torch.compile optimizations...")
+                self.model = torch.compile(self.model, mode="reduce-overhead")
+                self.prompt_generator = torch.compile(self.prompt_generator, mode="reduce-overhead")
+                print("✅ Torch compile applied successfully!")
+        except Exception as e:
+            print(f"⚠️ Torch compile not available or failed: {e}")
+
+        # Pre-warm GPU
+        try:
+            if self.device == "cuda":
+                dummy_input = torch.randn(1, 1024, dtype=torch.float16, device=self.device)
+                _ = self.sbert_model.encode(["test"], convert_to_tensor=True, device=self.device)
+                del dummy_input
+                torch.cuda.empty_cache()
+                print("✅ GPU pre-warmed successfully!")
+        except Exception as e:
+            print(f"⚠️ GPU pre-warming failed: {e}")
+
+    def create_cluster_embedding(self, pmid_abstracts: List[str], keywords: List[str]) -> torch.Tensor:
+        """
+        ENHANCED: Create better cluster embedding with keyword weighting
+        """
+        # Store for context
+        self._last_keywords = keywords
+        self._last_abstracts = pmid_abstracts
+
+        # Combine all abstracts
+        combined_abstracts = " ".join([clean_text(abstract) for abstract in pmid_abstracts if abstract])
+
+        # ENHANCED: Better keyword processing with importance weighting
+        if keywords:
+            clean_keywords = []
+            keyword_weights = []
+
+            for i, kw in enumerate(keywords):
+                if isinstance(kw, str):
+                    clean_kw = re.sub(r'\s*\([^)]+\)', '', kw).strip()
+                    if clean_kw and len(clean_kw) > 1:
+                        clean_keywords.append(clean_kw)
+                        # Higher weight for earlier keywords (assumed more important)
+                        keyword_weights.append(1.0 / (i + 1))
+
+            # Limit keywords but keep weights proportional
+            if len(clean_keywords) > 20:
+                clean_keywords = clean_keywords[:20]
+                keyword_weights = keyword_weights[:20]
+
+            # Normalize weights
+            if keyword_weights:
+                total_weight = sum(keyword_weights)
+                keyword_weights = [w/total_weight for w in keyword_weights]
+
+            # ENHANCED: Create weighted keyword text
+            keyword_text = ', '.join(clean_keywords)
+
+            # ENHANCED: Combine with emphasis on important keywords
+            important_keywords = clean_keywords[:5] if len(clean_keywords) >= 5 else clean_keywords
+            combined_text = f"{combined_abstracts}\n\nKey research topics: {', '.join(important_keywords)}. Additional concepts: {keyword_text}"
+        else:
+            combined_text = combined_abstracts
+
+        # Generate embedding with enhanced method
+        return self._compute_enhanced_embedding(combined_text, keywords)
+
+    def _compute_enhanced_embedding(self, text: str, keywords: List[str] = None) -> torch.Tensor:
+        """
+        ENHANCED: Compute embedding with better chunking and keyword integration
+        """
+        with torch.no_grad():
+            # Get main text embedding
+            text_embedding = self._compute_robust_embedding(text)
+
+            # ENHANCED: Add keyword embedding if available
+            if keywords and len(keywords) > 0:
+                # Create keyword-only embedding
+                keyword_text = ' [SEP] '.join(keywords[:15])  # Use separator tokens
+                keyword_embedding = self.sbert_model.encode(
+                    [keyword_text],
+                    convert_to_tensor=True,
+                    device=self.device,
+                    normalize_embeddings=True
+                ).squeeze(0).cpu()
+
+                # ENHANCED: Weighted combination (80% text, 20% keywords)
+                alpha = 0.85  # Text weight
+                beta = 0.15   # Keyword weight
+
+                combined_embedding = alpha * text_embedding + beta * keyword_embedding
+                combined_embedding = torch.nn.functional.normalize(combined_embedding.unsqueeze(0), p=2, dim=-1).squeeze(0)
+
+                return combined_embedding
+
+            return text_embedding
+
+    def _compute_robust_embedding(self, text: str) -> torch.Tensor:
+        """Compute robust embedding with chunking - optimized version"""
+        with torch.no_grad():
+            tokenized = self.sbert_model.tokenizer.encode(text, add_special_tokens=False)
+            total_tokens = len(tokenized)
+
+            if total_tokens <= 512:
+                embedding = self.sbert_model.encode(
+                    [text],
+                    convert_to_tensor=True,
+                    device=self.device,
+                    batch_size=1,
+                    show_progress_bar=False,
+                    normalize_embeddings=True
+                )
+            else:
+                # ENHANCED: Better chunking with overlap
+                chunks = []
+                chunk_weights = []
+
+                # Use sliding window with overlap
+                window_size = 512
+                stride = 256  # 50% overlap for better context
+
+                for i in range(0, total_tokens, stride):
+                    chunk_tokens = tokenized[i:i + window_size]
+                    if len(chunk_tokens) < 100:  # Skip tiny chunks
+                        break
+
+                    chunk_text = self.sbert_model.tokenizer.decode(chunk_tokens, skip_special_tokens=True)
+                    chunks.append(chunk_text)
+
+                    # ENHANCED: Position-based weighting (first and last chunks more important)
+                    position_weight = 1.2 if i == 0 else (1.1 if i + window_size >= total_tokens else 1.0)
+                    chunk_weights.append(position_weight * len(chunk_tokens))
+
+                # Process chunks in batches
+                chunk_batch_size = 16
+                chunk_embeddings_list = []
+
+                for i in range(0, len(chunks), chunk_batch_size):
+                    batch_chunks = chunks[i:i+chunk_batch_size]
+                    batch_embeds = self.sbert_model.encode(
+                        batch_chunks,
+                        convert_to_tensor=True,
+                        device=self.device,
+                        batch_size=len(batch_chunks),
+                        show_progress_bar=False,
+                        normalize_embeddings=True
+                    )
+                    chunk_embeddings_list.append(batch_embeds)
+
+                chunk_embeddings = torch.cat(chunk_embeddings_list, dim=0)
+                chunk_weights_tensor = torch.tensor(chunk_weights, dtype=torch.float16, device=chunk_embeddings.device)
+
+                # Normalize weights
+                chunk_weights_tensor = chunk_weights_tensor / chunk_weights_tensor.sum()
+
+                # Weighted average
+                embedding = torch.sum(chunk_embeddings * chunk_weights_tensor.unsqueeze(1), dim=0, keepdim=True)
+
+            return embedding.squeeze(0).cpu()
+
+    def generate_research_analysis(self, embedding: torch.Tensor, max_length: int = 500) -> Tuple[str, str, str]:
+        """
+        FIXED: Generate with corrected generation parameters
+        """
+        self.model.eval()
+        self.prompt_generator.eval()
+
+        # FIXED: Use compatible generation configurations
+        generation_configs = [
+            {
+                'name': 'high_quality',
+                'temperature': 0.7,
+                'top_p': 0.9,
+                'top_k': 50,
+                'num_beams': 5,
+                'do_sample': True,
+                'repetition_penalty': 1.15
+            },
+            {
+                'name': 'diverse_beam',
+                'num_beams': 5,
+                'num_beam_groups': 5,
+                'diversity_penalty': 0.5,
+                'do_sample': False,  # FIXED: Must be False for diverse beam search
+                'temperature': 1.0,  # Not used when do_sample=False
+                'repetition_penalty': 1.2
+            },
+            {
+                'name': 'focused',
+                'temperature': 0.6,
+                'top_p': 0.85,
+                'top_k': 40,
+                'num_beams': 6,
+                'do_sample': True,
+                'repetition_penalty': 1.1
+            }
+        ]
+
+        with torch.no_grad():
+            if embedding.dim() == 1:
+                embedding = embedding.unsqueeze(0)
+
+            embedding = embedding.to(self.device, dtype=torch.float16)
+            prefix_embeds = self.prompt_generator(embedding)
+
+            # ENHANCED: Better keyword context
+            if self._last_keywords:
+                # Clean keywords for better prompting
+                clean_keywords = []
+                for kw in self._last_keywords[:5]:
+                    clean_kw = re.sub(r'[_-]', ' ', str(kw)).strip()
+                    if clean_kw:
+                        clean_keywords.append(clean_kw)
+                keywords_text = ', '.join(clean_keywords) if clean_keywords else 'research topics'
+            else:
+                keywords_text = 'research topics'
+
+            # ENHANCED: Diverse vocabulary instruction prompt to reduce repetition
+            instruction_start = f"""<|begin_of_text|><|start_header_id|>system<|end_header_id|>
+
+You are a scientific theme analyst. Generate exactly three outputs for a biomedical topic:
+
+TITLE: [8-12 word distinctive title using diverse vocabulary - avoid repeating 'research', 'analysis', 'study'. Use terms like: mechanisms, pathways, connections, interactions, dynamics, networks, insights, perspectives, implications, applications]
+SHORT_SUMMARY: [2-3 sentences, 50-100 words describing the scientific domain and scope]
+ABSTRACT: [4-6 sentences, 150-300 words detailed description of mechanisms, pathways, and clinical significance]
+
+Use varied scientific terminology. Avoid repetitive language patterns. Focus on biological mechanisms, molecular pathways, clinical implications, and therapeutic potential.<|eot_id|><|start_header_id|>user<|end_header_id|>
+
+Generate content for biomedical domain involving: {keywords_text}<|eot_id|><|start_header_id|>assistant<|end_header_id|>
+
+TITLE: """
+
+            instruction_tokens = self.tokenizer(
+                instruction_start,
+                return_tensors="pt",
+                add_special_tokens=False
+            )
+            instruction_embeds = self.model.get_input_embeddings()(instruction_tokens["input_ids"].to(prefix_embeds.device))
+
+            full_inputs_embeds = torch.cat([prefix_embeds, instruction_embeds], dim=1)
+
+            seq_len = full_inputs_embeds.shape[1]
+            attention_mask = torch.ones((1, seq_len), dtype=torch.long, device=prefix_embeds.device)
+
+            # Try different generation strategies
+            generated_text = None
+            for config in generation_configs[:2]:  # Try first two configs
+                try:
+                    # Build generation kwargs based on config
+                    gen_kwargs = {
+                        'inputs_embeds': full_inputs_embeds,
+                        'attention_mask': attention_mask,
+                        'max_new_tokens': max_length,
+                        'min_new_tokens': 200,
+                        'num_beams': config.get('num_beams', 4),
+                        'no_repeat_ngram_size': 4,
+                        'length_penalty': 1.0,
+                        'early_stopping': False,
+                        'pad_token_id': self.tokenizer.pad_token_id,
+                        'eos_token_id': self.tokenizer.eos_token_id,
+                        'use_cache': True,
+                        'repetition_penalty': config.get('repetition_penalty', 1.1)
+                    }
+
+                    # Add config-specific parameters
+                    if 'num_beam_groups' in config:
+                        gen_kwargs['num_beam_groups'] = config['num_beam_groups']
+                    if 'diversity_penalty' in config:
+                        gen_kwargs['diversity_penalty'] = config['diversity_penalty']
+                    if 'do_sample' in config:
+                        gen_kwargs['do_sample'] = config['do_sample']
+                    if config.get('do_sample', False):  # Only add these if sampling
+                        gen_kwargs['temperature'] = config.get('temperature', 0.7)
+                        gen_kwargs['top_p'] = config.get('top_p', 0.9)
+                        gen_kwargs['top_k'] = config.get('top_k', 50)
+
+                    generated_ids = self.model.generate(**gen_kwargs)
+                    generated_text = self.tokenizer.decode(generated_ids[0], skip_special_tokens=True)
+
+                    # Extract generated part
+                    if "TITLE:" in generated_text:
+                        parts = generated_text.split("TITLE:")
+                        if len(parts) > 1:
+                            generated_text = "TITLE:" + parts[-1]
+
+                    # If we got a good generation, break
+                    if generated_text and len(generated_text) > 100:
+                        break
+
+                except Exception as e:
+                    if "diversity_penalty" not in str(e):  # Only print unexpected errors
+                        print(f"⚠️ Generation with {config['name']} config failed: {e}")
+                    continue
+
+        # Parse the output
+        if generated_text:
+            return self._parse_generated_output_enhanced(generated_text)
+        else:
+            # Fallback if all attempts failed
+            return self._generate_contextual_abstract(), self._generate_contextual_overview(), self._generate_contextual_title()
+
+    def _parse_generated_output_enhanced(self, text: str) -> Tuple[str, str, str]:
+        """
+        ENHANCED: Better parsing with validation and correction
+        """
+        text = text.strip()
+
+        # Clean up artifacts
+        text = re.sub(r'<\|.*?\|>', '', text).strip()
+
+        # ENHANCED: More robust regex patterns matching training format
+        title_match = re.search(
+            r'(?:TITLE|Title):?\s*([^\n]+?)(?=\n|SHORT_SUMMARY:|SHORT SUMMARY:|$)',
+            text,
+            re.IGNORECASE
+        )
+
+        short_match = re.search(
+            r'(?:SHORT[_ ]SUMMARY):?\s*([^\n]+(?:\n[^\n:]+)*?)(?=\nABSTRACT:|$)',
+            text,
+            re.IGNORECASE | re.DOTALL
+        )
+
+        abstract_match = re.search(
+            r'(?:ABSTRACT|Abstract):?\s*(.+?)(?=$)',
+            text,
+            re.IGNORECASE | re.DOTALL
+        )
+
+        title = title_match.group(1).strip() if title_match else ""
+        overview = short_match.group(1).strip() if short_match else ""
+        abstract = abstract_match.group(1).strip() if abstract_match else ""
+
+        # ENHANCED: Better validation and correction
+        title = self._validate_and_correct_title(title)
+        overview = self._validate_and_correct_overview(overview)
+        abstract = self._validate_and_correct_abstract(abstract)
+
+        # Final quality check
+        if not self._is_quality_output(title, overview, abstract):
+            # Try to salvage what we can
+            if not title:
+                title = self._generate_contextual_title()
+            if not overview:
+                overview = self._generate_contextual_overview()
+            if not abstract:
+                abstract = self._generate_contextual_abstract()
+
+        return abstract, overview, title
+
+    def _validate_and_correct_title(self, title: str) -> str:
+        """ENHANCED: Validate and correct title, removing repetitive patterns and repeated words"""
+        if not title:
+            return ""
+
+        # Remove common prefixes and suffixes
+        title = re.sub(r'^(TITLE:?\s*|Title:?\s*)', '', title, flags=re.IGNORECASE)
+        title = re.sub(r'^(Investigation of|Analysis of|Study of|Research on|Examination of)\s+', '', title, flags=re.IGNORECASE)
+
+        # ENHANCED: Remove more repetitive endings and patterns
+        repetitive_endings = [
+            r'\s+in Clinical Research Applications?$',
+            r'\s+in Biomedical Research$',
+            r'\s+in Healthcare Settings?$',
+            r'\s+in Medical Research$',
+            r'\s+Research Applications?$',
+            r'\s+Clinical Applications?$',
+            r'\s+Research Theme$',
+            r'\s+Theme Analysis$',
+            r'\s+Research Analysis$',
+            r'\s+Clinical Analysis$'
+        ]
+
+        for pattern in repetitive_endings:
+            title = re.sub(pattern, '', title, flags=re.IGNORECASE)
+
+        # ENHANCED: Remove repeated words within the title
+        title = self._remove_repeated_words(title)
+
+        # Clean whitespace
+        title = re.sub(r'\s+', ' ', title).strip()
+
+        # Enforce word count (8-15 words for more concise titles)
+        words = title.split()
+        if len(words) > 15:
+            # Find natural break point
+            for i in range(12, min(16, len(words))):
+                if words[i].lower() in ['and', 'with', 'through', 'via', 'using', 'from', 'to', 'in', 'for']:
+                    words = words[:i]
+                    break
+            else:
+                words = words[:15]
+            title = ' '.join(words)
+
+        # Ensure minimum length
+        if len(words) < 5:
+            return ""
+
+        # ENHANCED: Check for overused terms and suggest alternatives
+        title = self._avoid_overused_terms(title)
+
+        # Track word usage for future titles
+        self._track_title_words(title)
+
+        # Capitalize appropriately
+        return self._smart_capitalize(title)
+
+    def _remove_repeated_words(self, text: str) -> str:
+        """Remove repeated words within a title while preserving meaning"""
+        words = text.split()
+        if len(words) <= 3:
+            return text
+            
+        # Track word usage (case-insensitive)
+        seen_words = set()
+        filtered_words = []
+        
+        # Common words that can appear multiple times
+        allowed_repeats = {'and', 'or', 'of', 'in', 'for', 'with', 'the', 'a', 'an', 'to', 'from', 'by'}
+        
+        for word in words:
+            word_lower = word.lower()
+            # Allow common words to repeat, but remove other repetitions
+            if word_lower not in seen_words or word_lower in allowed_repeats:
+                filtered_words.append(word)
+                seen_words.add(word_lower)
+            # Special case: if removing this word would make title too short, keep it
+            elif len(filtered_words) < 6:
+                filtered_words.append(word)
+                
+        return ' '.join(filtered_words)
+
+    def _track_title_words(self, title: str) -> None:
+        """Track word usage across all generated titles"""
+        words = title.lower().split()
+        # Filter out common words that don't affect diversity
+        meaningful_words = [w for w in words if w not in {'and', 'or', 'of', 'in', 'for', 'with', 'the', 'a', 'an', 'to', 'from', 'by', 'on', 'at'}]
+        self._title_word_frequency.update(meaningful_words)
+
+    def _avoid_overused_terms(self, title: str) -> str:
+        """Replace overused terms with alternatives to improve diversity"""
+        words = title.split()
+        
+        # Replacement dictionary for overused terms
+        replacements = {
+            'research': ['investigation', 'exploration', 'inquiry', 'analysis'],
+            'analysis': ['examination', 'evaluation', 'assessment', 'investigation'],
+            'study': ['investigation', 'exploration', 'examination', 'inquiry'],
+            'application': ['implementation', 'utilization', 'deployment', 'use'],
+            'approach': ['strategy', 'method', 'technique', 'framework'],
+            'system': ['network', 'framework', 'mechanism', 'pathway'],
+            'method': ['technique', 'approach', 'strategy', 'protocol'],
+            'role': ['function', 'impact', 'influence', 'effect'],
+            'effect': ['impact', 'influence', 'consequence', 'outcome'],
+            'factor': ['element', 'component', 'determinant', 'variable']
+        }
+        
+        # Check each word for overuse
+        for i, word in enumerate(words):
+            word_lower = word.lower()
+            # If word is overused (appears more than 5 times) and has replacements
+            if (self._title_word_frequency[word_lower] > 5 and 
+                word_lower in replacements):
+                # Choose replacement based on current frequency
+                alternatives = replacements[word_lower]
+                best_alt = min(alternatives, key=lambda x: self._title_word_frequency[x])
+                # Only replace if the alternative is less used
+                if self._title_word_frequency[best_alt] < self._title_word_frequency[word_lower]:
+                    # Preserve original capitalization
+                    if word[0].isupper():
+                        words[i] = best_alt.capitalize()
+                    else:
+                        words[i] = best_alt
+        
+        return ' '.join(words)
+
+    def _validate_and_correct_overview(self, overview: str) -> str:
+        """ENHANCED: Validate and correct overview"""
+        if not overview:
+            return ""
+
+        # Remove label
+        overview = re.sub(r'^(SHORT[_ ]SUMMARY|OVERVIEW):?\s*', '', overview, flags=re.IGNORECASE)
+        overview = re.sub(r'\s+', ' ', overview).strip()
+
+        # Check length (should be 50-150 words)
+        words = overview.split()
+        if len(words) < 20 or len(words) > 150:
+            return ""
+
+        # Ensure it ends with proper punctuation
+        if overview and overview[-1] not in '.!?':
+            overview += '.'
+
+        return overview
+
+    def _validate_and_correct_abstract(self, abstract: str) -> str:
+        """ENHANCED: Validate and correct abstract"""
+        if not abstract:
+            return ""
+
+        # Remove label
+        abstract = re.sub(r'^(ABSTRACT):?\s*', '', abstract, flags=re.IGNORECASE)
+        abstract = re.sub(r'\s+', ' ', abstract).strip()
+
+        # Check length (should be 150-400 words)
+        words = abstract.split()
+        if len(words) < 50:
+            return ""
+
+        # Truncate if too long
+        if len(words) > 400:
+            # Try to find sentence boundary
+            sentences = re.split(r'(?<=[.!?])\s+', abstract)
+            result = []
+            word_count = 0
+            for sentence in sentences:
+                sentence_words = len(sentence.split())
+                if word_count + sentence_words <= 380:
+                    result.append(sentence)
+                    word_count += sentence_words
+                else:
+                    break
+            abstract = ' '.join(result)
+
+        # Ensure proper ending
+        if abstract and abstract[-1] not in '.!?':
+            abstract += '.'
+
+        return abstract
+
+    def _is_quality_output(self, title: str, overview: str, abstract: str) -> bool:
+        """Check if output meets quality standards"""
+        return (
+            len(title.split()) >= 5 and len(title.split()) <= 20 and
+            len(overview.split()) >= 20 and len(overview.split()) <= 150 and
+            len(abstract.split()) >= 50 and len(abstract.split()) <= 400 and
+            title != overview and title != abstract and overview != abstract
+        )
+
+    def _smart_capitalize(self, text: str) -> str:
+        """Smart capitalization for titles"""
+        words = text.split()
+        if not words:
+            return text
+
+        # Always capitalize first word
+        words[0] = words[0][0].upper() + words[0][1:] if len(words[0]) > 1 else words[0].upper()
+
+        # Small words that shouldn't be capitalized (unless first)
+        small_words = {'of', 'in', 'and', 'or', 'the', 'a', 'an', 'to', 'for', 'with', 'from', 'by', 'on', 'at'}
+
+        for i in range(1, len(words)):
+            if words[i].lower() not in small_words or i == len(words) - 1:
+                # Keep acronyms as is
+                if not words[i].isupper() or len(words[i]) > 4:
+                    words[i] = words[i][0].upper() + words[i][1:] if len(words[i]) > 1 else words[i].upper()
+
+        return ' '.join(words)
+
+    def _generate_contextual_title(self) -> str:
+        """ENHANCED: Generate diverse theme titles with varied vocabulary"""
+        if self._last_keywords and len(self._last_keywords) >= 2:
+            # Clean keywords
+            kw1 = re.sub(r'[_-]', ' ', str(self._last_keywords[0])).strip().title()
+            kw2 = re.sub(r'[_-]', ' ', str(self._last_keywords[1])).strip().title()
+
+            # ENHANCED: More diverse templates with varied vocabulary
+            templates = [
+                f"{kw1} and {kw2} Integration",
+                f"{kw1}-{kw2} Connections",
+                f"{kw1} Influences on {kw2}",
+                f"{kw2} Mechanisms in {kw1}",
+                f"{kw1} and {kw2}: Clinical Insights",
+                f"{kw1}-{kw2} Therapeutic Pathways",
+                f"{kw1} Interactions with {kw2}",
+                f"{kw2}-Mediated {kw1} Effects",
+                f"{kw1} and {kw2}: Biomedical Perspectives",
+                f"{kw1}-{kw2} Molecular Networks",
+                f"{kw1} Impact on {kw2} Regulation",
+                f"{kw2} Dynamics in {kw1} Context",
+                f"{kw1} and {kw2}: Translational Science",
+                f"{kw1}-{kw2} Disease Mechanisms",
+                f"{kw1} and {kw2}: Precision Medicine",
+                f"{kw2}-Associated {kw1} Pathways"
+            ]
+
+            # Select based on hash for consistency, but avoid repeating
+            base_hash = hash(''.join(self._last_keywords[:2]))
+
+            # Try to avoid recently used patterns
+            for i in range(len(templates)):
+                idx = (base_hash + i) % len(templates)
+                candidate = templates[idx]
+                pattern_key = f"{kw1[:3]}_{kw2[:3]}"  # Simple key for tracking
+
+                if self._title_patterns_used[pattern_key] < 3:  # Allow each pattern 3 times max
+                    self._title_patterns_used[pattern_key] += 1
+                    return candidate
+
+            # Fallback if all patterns used
+            return templates[base_hash % len(templates)]
+
+        return "Biomedical Mechanisms and Clinical Applications"
+
+    def _generate_contextual_overview(self) -> str:
+        """UPDATED: Generate theme overview using 'research theme covers' language"""
+        if self._last_keywords and len(self._last_keywords) >= 2:
+            # Clean keywords for natural language
+            clean_kw = []
+            for kw in self._last_keywords[:3]:
+                clean = re.sub(r'[_-]', ' ', str(kw)).strip().lower()
+                if clean:
+                    clean_kw.append(clean)
+
+            if len(clean_kw) >= 2:
+                return (f"This research theme covers the relationships between {clean_kw[0]} and {clean_kw[1]}, "
+                       f"encompassing significant implications for clinical practice. The theme covers "
+                       f"novel mechanisms that could lead to improved therapeutic strategies and patient outcomes.")
+
+        return ("This research theme covers important biomedical mechanisms with "
+                "significant clinical implications. The theme encompasses new insights for "
+                "developing more effective treatment strategies and improving patient care.")
+
+    def _generate_contextual_abstract(self) -> str:
+        """UPDATED: Generate theme abstract using theme-oriented language"""
+        if self._last_keywords and len(self._last_keywords) >= 3:
+            # Clean keywords
+            kw1 = re.sub(r'[_-]', ' ', str(self._last_keywords[0])).strip().lower()
+            kw2 = re.sub(r'[_-]', ' ', str(self._last_keywords[1])).strip().lower()
+            kw3 = re.sub(r'[_-]', ' ', str(self._last_keywords[2])).strip().lower()
+
+            return (f"This research theme covers the complex relationships between {kw1} and {kw2} "
+                   f"through comprehensive analysis of clinical and experimental data. The theme encompasses "
+                   f"novel interactions involving {kw3} that contribute to disease mechanisms and therapeutic responses. "
+                   f"This research theme covers previously unrecognized pathways that regulate these processes in clinical "
+                   f"populations. The theme demonstrates significant associations between these "
+                   f"factors and patient outcomes, with important implications for treatment selection "
+                   f"and optimization. This research theme provides a foundation for developing targeted "
+                   f"interventions and improving clinical care through personalized medicine approaches.")
+
+        return self._generate_fallback_abstract()
+
+    def _generate_fallback_title(self) -> str:
+        """ENHANCED: Generate diverse fallback titles"""
+        if self._last_keywords and len(self._last_keywords) >= 2:
+            kw1 = re.sub(r'[_-]', ' ', str(self._last_keywords[0])).strip().title()
+            kw2 = re.sub(r'[_-]', ' ', str(self._last_keywords[1])).strip().title()
+            fallback_patterns = [
+                f"{kw1} and {kw2}: Molecular Insights",
+                f"{kw1}-{kw2} Therapeutic Connections",
+                f"{kw1} Interactions with {kw2}",
+                f"{kw2}-Mediated {kw1} Pathways"
+            ]
+            # Use hash for consistent but varied selection
+            idx = hash(''.join(self._last_keywords[:2])) % len(fallback_patterns)
+            return fallback_patterns[idx]
+        return "Biomedical Mechanisms and Clinical Applications"
+
+    def _generate_fallback_overview(self) -> str:
+        """UPDATED: Generate fallback theme overview"""
+        return ("This research theme covers important insights into biomedical mechanisms "
+                "and their clinical applications. The theme encompasses significant implications "
+                "for improving patient care and developing new treatment strategies.")
+
+    def _generate_fallback_abstract(self) -> str:
+        """UPDATED: Generate fallback theme abstract"""
+        return ("This research theme covers complex biomedical mechanisms "
+                "through systematic analysis of clinical and experimental data. The theme encompasses "
+                "novel pathways and interactions that contribute to disease progression and treatment response. "
+                "This research theme covers important regulatory mechanisms that were previously unrecognized in clinical "
+                "populations. The theme has significant implications for developing "
+                "more effective therapeutic strategies and improving patient outcomes through "
+                "personalized medicine approaches. This research theme provides a foundation for future "
+                "research and clinical applications in precision medicine.")
+
+    # Memory management utilities
+    def cleanup_memory(self):
+        """Aggressive memory cleanup for long-running inference"""
+        torch.cuda.empty_cache()
+        gc.collect()
+        print("🧹 Memory cleanup completed")
+
+    def get_memory_stats(self):
+        """Get current GPU memory usage"""
+        if torch.cuda.is_available():
+            allocated = torch.cuda.memory_allocated() / 1024**3
+            reserved = torch.cuda.memory_reserved() / 1024**3
+            return f"GPU Memory - Allocated: {allocated:.2f}GB, Reserved: {reserved:.2f}GB"
+        return "CUDA not available"
+
+    def process_pickle_data(self, pickle_file_path: str, keywords_dict: Dict = None) -> List[Dict]:
+        """Process pickle file data with enhanced generation"""
+        print(f"📂 Loading data from {pickle_file_path}")
+
+        with open(pickle_file_path, 'rb') as f:
+            data = pickle.load(f)
+
+        results = []
+        num_clusters = data['metadata']['num_clusters']
+
+        print(f"🔄 Processing {num_clusters} clusters with enhanced generation...")
+
+        # Pre-allocate result list
+        results = [None] * num_clusters
+
+        # Process with progress bar
+        for cluster_idx in tqdm(range(num_clusters), desc="Generating analyses"):
+            try:
+                # Extract cluster data
+                cluster_docs = data['cluster_docs'][cluster_idx] if cluster_idx < len(data['cluster_docs']) else []
+                pmid_abstracts = data['pmid_abstracts'][cluster_idx] if cluster_idx < len(data['pmid_abstracts']) else []
+                keywords = keywords_dict.get(cluster_idx, []) if keywords_dict else []
+
+                # Create embedding with enhanced method
+                embedding = self.create_cluster_embedding(pmid_abstracts, keywords)
+
+                # Generate content with enhanced parameters
+                abstract, overview, title = self.generate_research_analysis(embedding, max_length=500)
+
+                results[cluster_idx] = {
+                    'cluster_id': cluster_idx,
+                    'abstract': abstract,
+                    'overview': overview,
+                    'title': title,
+                    'num_pmids': len(pmid_abstracts),
+                    'keywords': keywords[:10]
+                }
+
+                # Memory cleanup every 10 clusters
+                if cluster_idx % 10 == 0:
+                    torch.cuda.empty_cache()
+                    gc.collect()
+
+            except Exception as e:
+                print(f"⚠️ Error processing cluster {cluster_idx}: {e}")
+                results[cluster_idx] = {
+                    'cluster_id': cluster_idx,
+                    'abstract': self._generate_fallback_abstract(),
+                    'overview': self._generate_fallback_overview(),
+                    'title': f"Research Theme {cluster_idx} Analysis",
+                    'num_pmids': 0,
+                    'keywords': []
+                }
+
+        # Final cleanup
+        torch.cuda.empty_cache()
+        gc.collect()
+
+        # Filter out None results
+        results = [r for r in results if r is not None]
+
+        return results
+
+    def save_results_tsv(self, results: List[Dict], output_path: str = None, prefix: str = "research_analyses"):
+        """Save results to timestamped TSV file"""
+        if output_path is None:
+            timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
+            output_path = f"{prefix}_{timestamp}.tsv"
+
+        df = pd.DataFrame(results)
+        df.to_csv(output_path, sep='\t', index=False)
+        print(f"💾 Results saved to: {output_path}")
+        return output_path
+
+    # Backward compatibility wrapper
+    def generate_research_summary(self, embedding: torch.Tensor, max_length: int = 500) -> Tuple[str, str, str]:
+        """Backward compatibility wrapper"""
+        return self.generate_research_analysis(embedding, max_length)
+
+# Convenience function for easy usage
+def load_model_and_generate(model_dir: str, pickle_files: List[str], keywords_dict: Dict = None,
+                           output_prefix: str = "research_analyses") -> List[str]:
+    """
+    Convenience function to load model and generate analyses for multiple pickle files
+    """
+    print("🚀 Initializing model with fixed generation parameters...")
+    model = ScientificModelInference(model_dir)
+
+    print(f"📊 {model.get_memory_stats()}")
+
+    output_files = []
+
+    for i, pickle_file in enumerate(pickle_files):
+        print(f"\n📋 Processing {pickle_file} ({i+1}/{len(pickle_files)})")
+
+        # Process data with enhanced generation
+        results = model.process_pickle_data(pickle_file, keywords_dict)
+
+        # Generate unique output name
+        period_name = Path(pickle_file).stem
+        output_path = model.save_results_tsv(results, prefix=f"{output_prefix}_{period_name}")
+        output_files.append(output_path)
+
+        # Memory cleanup between files
+        if len(pickle_files) > 1:
+            model.cleanup_memory()
+            print(f"📊 {model.get_memory_stats()}")
+
+    print(f"🎉 Completed processing {len(pickle_files)} files with improved titles!")
+    return output_files