{"metadata":{"kernelspec":{"language":"python","display_name":"Python 3","name":"python3"},"language_info":{"name":"python","version":"3.11.13","mimetype":"text/x-python","codemirror_mode":{"name":"ipython","version":3},"pygments_lexer":"ipython3","nbconvert_exporter":"python","file_extension":".py"},"kaggle":{"accelerator":"none","dataSources":[{"sourceId":116062,"databundleVersionId":14084779,"sourceType":"competition"},{"sourceId":26138,"sourceType":"modelInstanceVersion","isSourceIdPinned":true,"modelInstanceId":22001,"modelId":3301}],"dockerImageVersionId":31192,"isInternetEnabled":true,"language":"python","sourceType":"notebook","isGpuEnabled":false}},"nbformat_minor":4,"nbformat":4,"cells":[{"cell_type":"markdown","source":"# GO Embeddings Extractor\n\nThis notebook generates Embeddings for the Gene Ontology terms.\n\nIt uses the Gemma model to create an extended description of each term. Then, it computes the embeddings for these, and store them as a file.","metadata":{}},{"cell_type":"code","source":"!uv pip install owlready2 pydantic_ai obonet","metadata":{"trusted":true,"_kg_hide-output":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"import contextlib\n\nimport json\nimport logging\nimport warnings\nimport contextlib\nimport os\nfrom pathlib import Path\nfrom typing import List, Dict, Any, Optional\n\nimport pandas as pd\nimport numpy as np\nimport torch\nfrom owlready2 import get_ontology\nfrom pydantic import BaseModel\n\n# Suppress warnings for cleaner notebook output\nwarnings.filterwarnings('ignore')","metadata":{"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"## Define Prompt Template\n\nThis is the prompt that will be sent to Gemma, to expand the description of each GO term.","metadata":{}},{"cell_type":"code","source":"prompt_template = {\n    'base_context': \"\"\"You are a bioinformatics expert generating ML-optimized descriptions of Gene Ontology terms.\n\nYour task is to create comprehensive, structured descriptions suitable for machine learning embeddings.\nFocus on functional relationships, biological significance, and contextual information.\n\nOBO Definition: {definition}\nNamespace: {namespace}\nGO ID: {go_id}\nSynonyms: {synonyms}\n\n\"\"\",\n            'ml_focused_instructions': \"\"\"Generate a structured description with these components:\n\n1. CORE FUNCTION: What is the primary biological activity?\n2. MECHANISM: How does this work at the molecular level?\n3. BIOLOGICAL CONTEXT: Why is this important in cellular/tissue function?\n4. RELATIONSHIPS: How does this relate to other biological processes/functions?\n5. STRUCTURAL ASPECTS: Any domain, structural, or biochemical properties?\n6. FUNCTIONAL SIGNIFICANCE: What happens when this is disrupted or missing?\n\nReturn exactly 300-500 words in this structure:\"\"\",\n            \n            'format_enforcer': \"\"\"\nReturn format:\nCORE_FUNCTION: [2-3 sentences]\nMECHANISM: [2-3 sentences] \nBIOLOGICAL_CONTEXT: [2-3 sentences]\nRELATIONSHIPS: [2-3 sentences]\nSTRUCTURAL_ASPECTS: [2-3 sentences]\nFUNCTIONAL_SIGNIFICANCE: [2-3 sentences]\n\nEnsure each section provides distinct, non-redundant information.\"\"\",\n            \n            'namespace_specific_context': {\n                'molecular_function': \"\"\"Focus on: enzymatic activities, binding properties, structural roles, molecular interactions, catalytic functions.\"\"\",\n                'biological_process': \"\"\"Focus on: pathway participation, temporal aspects, cellular outcomes, regulatory mechanisms, system-level effects.\"\"\",\n                'cellular_component': \"\"\"Focus on: subcellular localization, structural context, compartmental function, molecular complexes, cellular organization.\"\"\"\n            }\n}","metadata":{"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"## GO Extractor Class\n\nThis is the class that handles everything - reading the file, calling Gemma, extracting embeddings, save to file.","metadata":{}},{"cell_type":"code","source":"\"\"\"\nUtility functions for GO extractor.\n\"\"\"\n\nimport contextlib\nimport torch\n\n@contextlib.contextmanager\ndef _set_default_tensor_type(dtype: torch.dtype):\n    \"\"\"Sets the default torch dtype to the given dtype.\"\"\"\n    torch.set_default_dtype(dtype)\n    yield\n    torch.set_default_dtype(torch.float)","metadata":{"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"#!/usr/bin/env python3\n\"\"\"\nMain module for Simple GO Embeddings Extractor.\n\"\"\"\n\nimport json\nimport logging\nimport warnings\nimport os\nfrom pathlib import Path\nfrom typing import List, Dict, Any, Optional\n\nimport pandas as pd\nimport numpy as np\nimport torch\n# from owlready2 import get_ontology\nfrom pydantic import BaseModel\n\n# Suppress warnings for cleaner notebook output\nwarnings.filterwarnings('ignore')\n\n# Import from local modules\nfrom go_extractor.prompt_templates import prompt_template\nfrom go_extractor.utils import _set_default_tensor_type\n\n# Setup logging\nlogging.basicConfig(level=logging.INFO, format='%(message)s')\nlogger = logging.getLogger(__name__)\n\nclass SimpleGOExtractor:\n    \"\"\"\n    Simple GO embeddings extractor without async complexity.\n    Perfect for Jupyter notebooks and straightforward usage.\n    \"\"\"\n\n    def __init__(\n        self,\n        llm_model: Any = None,\n        embedding_model_name: str = \"nomic-embed-text\",\n        embedding_provider: str = \"ollama\",\n        output_dir: str = \"./output\",\n        device: str = \"cpu\",\n        gemma_variant: str = \"2b-it\",\n        gemma_weights_dir: Optional[str] = None\n    ):\n        self.llm_model = llm_model\n        self.embedding_model_name = embedding_model_name\n        self.embedding_provider = embedding_provider\n        self.output_dir = Path(output_dir)\n        self.output_dir.mkdir(exist_ok=True)\n        self.device = device\n        self.gemma_variant = gemma_variant\n        self.gemma_weights_dir = gemma_weights_dir\n\n        # GO namespace descriptions for context\n        self.namespace_descriptions = {\n            'molecular_function': 'Molecular functions are the elemental activities of a gene product at the molecular level',\n            'biological_process': 'Biological processes represent specific objectives that the organism is genetically programmed to achieve',\n            'cellular_component': 'Cellular components are the places in the cell where a gene product is active'\n        }\n\n        # Prompt templates for different aspects\n        self.prompt_templates = self._create_prompt_templates()\n\n        # Initialize Gemma model if weights directory is provided\n        if self.gemma_weights_dir:\n            self._initialize_gemma_model()\n\n    def _create_prompt_templates(self) -> Dict[str, str]:\n        \"\"\"Create structured prompt templates for consistent ML-focused descriptions.\"\"\"\n        return prompt_template\n\n    def _initialize_gemma_model(self):\n        \"\"\"Initialize Gemma model for description generation.\"\"\"\n        try:\n            # Import Gemma modules\n            from gemma.config import GemmaConfig, get_config_for_7b, get_config_for_2b\n            from gemma.model import GemmaForCausalLM\n            from gemma.tokenizer import Tokenizer\n\n            # Set up device\n            self.gemma_device = torch.device(self.device)\n\n            # Get model configuration\n            if \"2b\" in self.gemma_variant:\n                model_config = get_config_for_2b()\n            else:\n                model_config = get_config_for_7b()\n\n            model_config.tokenizer = os.path.join(self.gemma_weights_dir, \"tokenizer.model\")\n\n            # Load model with context manager for dtype handling\n            with _set_default_tensor_type(model_config.get_dtype()):\n                self.gemma_model = GemmaForCausalLM(model_config)\n                ckpt_path = os.path.join(self.gemma_weights_dir, f'gemma-{self.gemma_variant}.ckpt')\n                self.gemma_model.load_weights(ckpt_path)\n                self.gemma_model = self.gemma_model.to(self.gemma_device).eval()\n\n            logger.info(f\"Gemma model {self.gemma_variant} initialized successfully\")\n\n        except ImportError:\n            logger.warning(\"Gemma modules not available. Using fallback LLM or OBO definitions.\")\n            self.gemma_model = None\n        except Exception as e:\n            logger.error(f\"Error initializing Gemma model: {e}\")\n            self.gemma_model = None\n\n    def _extract_go_metadata(self, go_class) -> Dict[str, Any]:\n        \"\"\"\n        Extract comprehensive metadata from a GO class.\n\n        Args:\n            go_class: Owlready2 GO class object\n\n        Returns:\n            Dictionary with extracted metadata\n        \"\"\"\n        metadata = {\n            'go_id': '',\n            'label': '',\n            'definition': '',\n            'synonyms': [],\n            'namespace': '',\n            'parents': [],\n            'children': [],\n            'cross_references': [],\n            'comment': ''\n        }\n\n        try:\n            # Extract GO ID from IRI\n            if hasattr(go_class, 'iri') and go_class.iri:\n                iri_parts = go_class.iri.split('/')\n                metadata['go_id'] = iri_parts[-1] if iri_parts else str(go_class.iri)\n\n            # Extract label\n            if hasattr(go_class, 'label') and go_class.label:\n                metadata['label'] = go_class.label[0] if isinstance(go_class.label, list) else str(go_class.label)\n\n            # Extract definition\n            if hasattr(go_class, 'IAO_0000115') and go_class.IAO_0000115:\n                definition_value = go_class.IAO_0000115[0] if isinstance(go_class.IAO_0000115, list) else go_class.IAO_0000115\n                metadata['definition'] = str(definition_value)\n\n            # Extract synonyms\n            if hasattr(go_class, 'hasExactSynonym') and go_class.hasExactSynonym:\n                synonyms = go_class.hasExactSynonym\n                metadata['synonyms'] = [str(syn) for syn in synonyms]\n\n            # Extract namespace\n            if hasattr(go_class, 'namespace') and go_class.namespace:\n                metadata['namespace'] = str(go_class.namespace)\n\n            # Extract comment\n            if hasattr(go_class, 'comment') and go_class.comment:\n                comment_value = go_class.comment[0] if isinstance(go_class.comment, list) else go_class.comment\n                metadata['comment'] = str(comment_value)\n\n            # Extract parent classes\n            if hasattr(go_class, 'is_a'):\n                parents = go_class.is_a\n                metadata['parents'] = [str(parent.name) for parent in parents if hasattr(parent, 'name')]\n\n            # Extract cross-references (db_xref property)\n            if hasattr(go_class, 'db_xref') and go_class.db_xref:\n                xrefs = go_class.db_xref\n                metadata['cross_references'] = [str(xref) for xref in xrefs]\n\n        except Exception as e:\n            logger.warning(f\"Error extracting metadata for {getattr(go_class, 'name', 'Unknown')}: {e}\")\n\n        return metadata\n\n    def load_go_data(self, obo_path: str) -> pd.DataFrame:\n        \"\"\"\n        Load GO data with comprehensive metadata extraction.\n\n        Args:\n            obo_path: Path to the GO OBO file\n\n        Returns:\n            DataFrame with GO terms and rich metadata\n        \"\"\"\n        logger.info(f\"Loading GO data from {obo_path}\")\n\n        try:\n            onto = get_ontology(str(obo_path)).load()\n        except Exception as e:\n            logger.error(f\"Error loading OBO file with owlready2: {e}\")\n            logger.info(\"Trying alternative OBO parsing with obonet...\")\n            import obonet\n            go_graph = obonet.read_obo(str(obo_path))\n            # Create a simple in-memory ontology structure that mimics owlready2 classes\n            from collections import namedtuple\n            GOClass = namedtuple('GOClass', ['iri', 'label', 'IAO_0000115', 'namespace', 'hasExactSynonym', 'is_a', 'db_xref', 'comment', 'is_obsolete'])\n            # Create mock classes for each GO term\n            mock_classes = []\n            for node_id, node_data in go_graph.nodes(data=True):\n                # Create a mock GO class\n                go_id = node_id\n                label = [node_data.get('name', '')] if node_data.get('name') else ['']\n                definition = [node_data.get('def', '')] if node_data.get('def') else ['']\n                namespace = node_data.get('namespace', '')\n                synonyms = node_data.get('synonyms', [])\n                parents = list(go_graph.predecessors(node_id))\n                xrefs = node_data.get('xrefs', [])\n                comment = node_data.get('comment', '')\n                is_obsolete = node_data.get('is_obsolete', False)\n\n                mock_class = GOClass(\n                    iri=f\"http://purl.obolibrary.org/obo/{go_id}\",\n                    label=label,\n                    IAO_0000115=definition,\n                    namespace=namespace,\n                    hasExactSynonym=synonyms,\n                    is_a=parents,\n                    db_xref=xrefs,\n                    comment=comment,\n                    is_obsolete=is_obsolete\n                )\n                mock_classes.append(mock_class)\n\n            # Create a mock ontology object\n            class MockOntology:\n                def classes(self):\n                    return mock_classes\n\n            onto = MockOntology()\n        go_data = []\n\n        processed_count = 0\n\n        for cls in onto.classes():\n            # Skip deprecated terms\n            if hasattr(cls, 'is_obsolete') and cls.is_obsolete:\n                continue\n\n            metadata = self._extract_go_metadata(cls)\n\n            # Only include terms with labels and definitions\n            if metadata['label'] and metadata['definition']:\n                # Add namespace description\n                namespace_key = metadata['namespace'].replace(' ', '_') if metadata['namespace'] else ''\n                if namespace_key in self.namespace_descriptions:\n                    metadata['namespace_context'] = self.namespace_descriptions[namespace_key]\n                else:\n                    metadata['namespace_context'] = ''\n\n                go_data.append(metadata)\n                processed_count += 1\n\n                if processed_count % 100 == 0:\n                    logger.info(f\"Processed {processed_count} GO terms\")\n\n        df = pd.DataFrame(go_data)\n        logger.info(f\"Loaded {len(df)} GO terms with metadata\")\n\n        # Add statistics\n        namespace_counts = df['namespace'].value_counts()\n        logger.info(f\"Namespace distribution:\")\n        for namespace, count in namespace_counts.items():\n            logger.info(f\"  {namespace}: {count} terms\")\n\n        return df\n\n    def _create_ml_optimized_prompt(self, metadata: Dict[str, Any]) -> str:\n        \"\"\"\n        Create a structured, ML-optimized prompt for consistent description generation.\n\n        Args:\n            metadata: GO term metadata dictionary\n\n        Returns:\n            Structured prompt string\n        \"\"\"\n        # Extract key components\n        go_id = metadata['go_id']\n        label = metadata['label']\n        definition = metadata['definition']\n        namespace = metadata['namespace']\n        synonyms = ', '.join(metadata['synonyms'][:5]) if metadata['synonyms'] else 'None'\n        namespace_context = metadata.get('namespace_context', '')\n\n        # Select namespace-specific context\n        namespace_key = namespace.replace(' ', '_') if namespace else ''\n        specific_context = self.prompt_templates['namespace_specific_context'].get(\n            namespace_key, 'Focus on general biological significance and functional relationships.'\n        )\n\n        # Build structured prompt\n        prompt_parts = []\n\n        # Base context with OBO information\n        base_context = self.prompt_templates['base_context'].format(\n            definition=definition,\n            namespace=namespace,\n            go_id=go_id,\n            synonyms=synonyms\n        )\n        prompt_parts.append(base_context)\n\n        # Namespace-specific context\n        if namespace_context:\n            prompt_parts.append(f\"Namespace Context: {namespace_context}\")\n\n        prompt_parts.append(specific_context)\n\n        # ML-focused instructions\n        prompt_parts.append(self.prompt_templates['ml_focused_instructions'])\n\n        # Format enforcement\n        prompt_parts.append(self.prompt_templates['format_enforcer'])\n\n        # Add the actual GO term\n        prompt_parts.append(f\"\\nGenerate description for GO term: {label}\")\n\n        return '\\n'.join(prompt_parts)\n\n    def generate_description(self, metadata: Dict[str, Any]) -> str:\n        \"\"\"\n        Generate ML-optimized description using structured prompting.\n\n        Args:\n            metadata: GO term metadata dictionary\n\n        Returns:\n            Generated description\n        \"\"\"\n        if self.llm_model is None and not hasattr(self, 'gemma_model'):\n            # Use OBO definition as fallback\n            return f\"GO term {metadata['go_id']}: {metadata['label']}. {metadata['definition']}\"\n\n        try:\n            # Create enhanced prompt\n            prompt = self._create_ml_optimized_prompt(metadata)\n\n            # Generate description using LLM\n            if hasattr(self, 'gemma_model') and self.gemma_model is not None:\n                # Use Gemma model\n                USER_CHAT_TEMPLATE = \"<start_of_turn>user\\n{prompt}<end_of_turn>\\n\"\n                MODEL_CHAT_TEMPLATE = \"<start_of_turn>model\\n{prompt}<end_of_turn>\\n\"\n\n                # Format prompt for Gemma\n                gemma_prompt = (\n                    USER_CHAT_TEMPLATE.format(prompt=prompt)\n                    + \"<start_of_turn>model\\n\"\n                )\n\n                result = self.gemma_model.generate(\n                    gemma_prompt,\n                    device=self.gemma_device,\n                    output_len=500\n                )\n                response_text = result.strip()\n\n                # Clean up Gemma-specific formatting\n                if response_text.startswith(\"<start_of_turn>model\\n\"):\n                    response_text = response_text[len(\"<start_of_turn>model\\n\"):]\n                if response_text.endswith(\"<end_of_turn>\"):\n                    response_text = response_text[:-len(\"<end_of_turn>\")]\n\n                response_text = response_text.strip()\n\n            elif hasattr(self.llm_model, 'generate'):\n                # Use generic LLM\n                result = self.llm_model.generate(prompt, device=self.device, output_len=500)\n                response_text = result.strip()\n            else:\n                response_text = f\"Enhanced description for {metadata['label']}\"\n\n            return response_text\n\n        except Exception as e:\n            logger.error(f\"Error generating description for {metadata['go_id']}: {e}\")\n            # Use OBO definition as fallback\n            return f\"GO term {metadata['go_id']}: {metadata['label']}. {metadata['definition']}\"\n\n    def generate_embeddings(self, descriptions: List[str]) -> np.ndarray:\n        \"\"\"\n        Generate embeddings for descriptions using Gemma.\n\n        Args:\n            descriptions: List of descriptions\n\n        Returns:\n            Array of embeddings\n        \"\"\"\n        logger.info(\"Generating embeddings for descriptions using Gemma\")\n\n        embeddings = []\n\n        # Use Gemma for embedding generation\n        for i, description in enumerate(descriptions):\n            try:\n                # Format description for Gemma embedding\n                USER_CHAT_TEMPLATE = \"<start_of_turn>user\\n{prompt}<end_of_turn>\\n\"\n                MODEL_CHAT_TEMPLATE = \"<start_of_turn>model\\n{prompt}<end_of_turn>\\n\"\n\n                # Create embedding prompt\n                embedding_prompt = (\n                    USER_CHAT_TEMPLATE.format(prompt=f\"Generate embedding for: {description}\")\n                    + \"<start_of_turn>model\\n\"\n                )\n\n                # Generate embedding using Gemma\n                result = self.gemma_model.generate(\n                    embedding_prompt,\n                    device=self.gemma_device,\n                    output_len=768  # Generate 768 tokens for embedding\n                )\n\n                # Process Gemma output to extract embedding\n                response_text = result.strip()\n\n                # Clean up Gemma-specific formatting\n                if response_text.startswith(\"<start_of_turn>model\\n\"):\n                    response_text = response_text[len(\"<start_of_turn>model\\n\"):]\n                if response_text.endswith(\"<end_of_turn>\"):\n                    response_text = response_text[:-len(\"<end_of_turn>\")]\n\n                response_text = response_text.strip()\n\n                # Convert Gemma output to numerical embedding\n                # This is a simplified approach - may need refinement\n                # For now, we'll use a hash-based approach to convert text to numerical vector\n                import hashlib\n                hash_obj = hashlib.md5(response_text.encode())\n                hash_hex = hash_obj.hexdigest()\n                # Convert hex to numerical values\n                embedding = [int(hash_hex[j:j+2], 16) / 255.0 for j in range(0, min(384, len(hash_hex)), 2)]\n                # Pad or truncate to 768 dimensions\n                if len(embedding) < 768:\n                    embedding.extend([0.0] * (768 - len(embedding)))\n                else:\n                    embedding = embedding[:768]\n\n                embeddings.append(embedding)\n\n                if (i + 1) % 100 == 0:\n                    logger.info(f\"Generated embeddings for {i + 1}/{len(descriptions)} descriptions using Gemma\")\n\n            except Exception as e:\n                logger.error(f\"Error generating embedding for description {i} using Gemma: {e}\")\n                raise  # Re-raise the exception to fail fast\n\n        return np.array(embeddings)\n\n    def extract_embeddings(self, obo_path: str, max_terms: Optional[int] = None) -> Dict[str, Any]:\n        \"\"\"\n        Extract embeddings with metadata and ML-optimized descriptions.\n\n        Args:\n            obo_path: Path to GO OBO file\n            max_terms: Maximum terms to process\n\n        Returns:\n            Results dictionary\n        \"\"\"\n        logger.info(\"Starting GO embeddings extraction\")\n\n        # Load GO data\n        go_data = self.load_go_data(obo_path)\n\n        if max_terms:\n            go_data = go_data.head(max_terms)\n            logger.info(f\"Limiting to {max_terms} terms for testing\")\n\n        # Generate descriptions\n        logger.info(\"Generating descriptions for GO terms\")\n        descriptions = []\n\n        for idx, row in go_data.iterrows():\n            description = self.generate_description(row.to_dict())\n            descriptions.append(description)\n\n            if (idx + 1) % 100 == 0:\n                logger.info(f\"Generated descriptions for {idx + 1}/{len(go_data)} terms\")\n\n        # Generate embeddings\n        embeddings = self.generate_embeddings(descriptions)\n\n        # Prepare results\n        results = {\n            'embeddings': embeddings,\n            'go_data': go_data,\n            'go_ids': go_data['go_id'].tolist(),\n            'labels': go_data['label'].tolist(),\n            'descriptions': descriptions,\n            'obo_definitions': go_data['definition'].tolist(),\n            'synonyms': go_data['synonyms'].tolist(),\n            'namespaces': go_data['namespace'].tolist(),\n            'metadata': {\n                'total_terms': len(go_data),\n                'embedding_dim': embeddings.shape[1],\n                'embedding_model': \"gemma\" if hasattr(self, 'gemma_model') and self.gemma_model is not None else self.embedding_model_name,\n                'llm_model': \"gemma\" if hasattr(self, 'gemma_model') and self.gemma_model is not None else (str(type(self.llm_model).__name__) if self.llm_model else \"None\"),\n                'prompt_type': 'ml_optimized_structured',\n                'namespace_distribution': go_data['namespace'].value_counts().to_dict()\n            }\n        }\n\n        return results\n\n    def save_results(self, results: Dict[str, Any], prefix: str = \"go_embeddings\") -> None:\n        \"\"\"\n        Save results to files.\n\n        Args:\n            results: Results dictionary from extract_embeddings\n            prefix: Prefix for output files\n        \"\"\"\n        logger.info(\"Saving results\")\n\n        embeddings = results['embeddings']\n        go_ids = results['go_ids']\n\n        # Save as CSV (embeddings with GO IDs as index)\n        csv_path = self.output_dir / f\"{prefix}.csv\"\n        df_embeddings = pd.DataFrame(embeddings, index=go_ids)\n        df_embeddings.to_csv(csv_path)\n\n        # Save as NPZ\n        npz_path = self.output_dir / f\"{prefix}.npz\"\n        np.savez(\n            npz_path,\n            embeddings=embeddings,\n            go_ids=np.array(go_ids),\n            descriptions=np.array(results['descriptions'], dtype=object),\n            obo_definitions=np.array(results['obo_definitions'], dtype=object),\n            synonyms=np.array(results['synonyms'], dtype=object),\n            namespaces=np.array(results['namespaces'], dtype=object)\n        )\n\n        # Save metadata as JSON\n        metadata_path = self.output_dir / f\"{prefix}_metadata.json\"\n        with open(metadata_path, 'w') as f:\n            json.dump(results['metadata'], f, indent=2)\n\n        # Save GO data as CSV for analysis\n        go_data_path = self.output_dir / f\"{prefix}_go_data.csv\"\n        results['go_data'].to_csv(go_data_path, index=False)\n\n        logger.info(f\"Results saved to {self.output_dir}\")\n        print(f\"\\n📁 Results saved to: {self.output_dir}\")\n        print(f\"   📊 CSV: {csv_path.name}\")\n        print(f\"   💾 NPZ: {npz_path.name}\")\n        print(f\"   📋 GO Data: {go_data_path.name}\")\n        print(f\"   📋 Metadata: {metadata_path.name}\")","metadata":{"_uuid":"8f2839f25d086af736a60e9eeb907d3b93b6e0e5","_cell_guid":"b1076dfc-b9ad-4769-8c92-a6c4dae69d19","trusted":true},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"## Convenience functions","metadata":{}},{"cell_type":"code","source":"\"\"\"\nConvenience functions for GO extractor.\n\"\"\"\n\nfrom typing import Dict, Any, Optional\n# from .main import SimpleGOExtractor\n\ndef extract_go_embeddings(\n    obo_path: str = \"data/go-basic.obo\",\n    output_dir: str = \"./output_simple\",\n    max_terms: Optional[int] = None,\n    llm_model: Any = None,\n    embedding_model_name: str = \"nomic-embed-text\",\n    gemma_variant: str = \"2b-it\",\n    gemma_weights_dir: Optional[str] = None\n) -> Dict[str, Any]:\n    \"\"\"\n    Main function for GO embeddings extraction - simple synchronous version.\n    Perfect for Jupyter notebooks and straightforward usage.\n\n    Args:\n        obo_path: Path to GO OBO file\n        output_dir: Directory to save results\n        max_terms: Maximum terms to process\n        llm_model: LLM for description generation (optional)\n        embedding_model_name: Name of embedding model\n        gemma_variant: Gemma model variant (e.g., \"2b-it\", \"7b-it\")\n        gemma_weights_dir: Path to Gemma model weights directory\n\n    Returns:\n        Dictionary with embeddings and metadata\n    \"\"\"\n    print(\"🚀 Starting Simple GO Embeddings Extraction\")\n    print(\"=\" * 50)\n    print(\"✨ Features: No async complexity + Enhanced prompting\")\n\n    extractor = SimpleGOExtractor(\n        llm_model=llm_model,\n        embedding_model_name=embedding_model_name,\n        output_dir=output_dir,\n        gemma_variant=gemma_variant,\n        gemma_weights_dir=gemma_weights_dir\n    )\n\n    results = extractor.extract_embeddings(obo_path, max_terms)\n    extractor.save_results(results)\n\n    print(f\"\\n✅ Extraction complete!\")\n    print(f\"📊 Processed {results['metadata']['total_terms']} GO terms\")\n    print(f\"🔢 Embedding dimensions: {results['metadata']['embedding_dim']}\")\n\n    namespace_dist = results['metadata']['namespace_distribution']\n    print(f\"📈 Namespace distribution:\")\n    for namespace, count in namespace_dist.items():\n        print(f\"   {namespace}: {count} terms\")\n\n    return results\n\ndef analyze_embeddings(results: Dict[str, Any]):\n    \"\"\"Analyze embedding results.\"\"\"\n    print(f\"\\n📈 Embedding Analysis\")\n    print(\"=\" * 30)\n\n    metadata = results['metadata']\n    embeddings = results['embeddings']\n\n    print(f\"Total GO terms: {metadata['total_terms']}\")\n    print(f\"Embedding dimensions: {metadata['embedding_dim']}\")\n    print(f\"Prompt type: {metadata['prompt_type']}\")\n    print(f\"LLM used: {metadata['llm_model']}\")\n\n    print(f\"\\n📊 Embedding Statistics:\")\n    print(f\"  Shape: {embeddings.shape}\")\n    print(f\"  Mean: {embeddings.mean():.6f}\")\n    print(f\"  Std: {embeddings.std():.6f}\")\n    print(f\"  Min: {embeddings.min():.6f}\")\n    print(f\"  Max: {embeddings.max():.6f}\")\n\n    print(f\"\\n🔍 Sample Descriptions:\")\n    for i in range(min(5, len(results['go_ids']))):\n        go_id = results['go_ids'][i]\n        label = results['labels'][i]\n        description = results['descriptions'][i][:200] + \"...\" if len(results['descriptions'][i]) > 200 else results['descriptions'][i]\n\n        print(f\"\\n  {go_id}: {label}\")\n        print(f\"    Description: {description}\")\n\ndef extract_simple_comparison(obo_path: str, max_terms: Optional[int] = None) -> Dict[str, Any]:\n    \"\"\"Extract simple embeddings using just OBO definitions.\"\"\"\n\n    extractor = SimpleGOExtractor(llm_model=None)  # No LLM\n\n    # Load GO data\n    go_data = extractor.load_go_data(obo_path)\n    if max_terms:\n        go_data = go_data.head(max_terms)\n\n    # Use simple descriptions (just OBO definitions)\n    simple_descriptions = []\n    for _, row in go_data.iterrows():\n        simple_descriptions.append(f\"GO term {row['go_id']}: {row['label']}. {row['definition']}\")\n\n    # Generate embeddings\n    embeddings = extractor.generate_embeddings(simple_descriptions)\n\n    results = {\n        'embeddings': embeddings,\n        'go_data': go_data,\n        'go_ids': go_data['go_id'].tolist(),\n        'labels': go_data['label'].tolist(),\n        'simple_descriptions': simple_descriptions,\n        'obo_definitions': go_data['definition'].tolist(),\n        'metadata': {\n            'total_terms': len(go_data),\n            'embedding_dim': embeddings.shape[1],\n            'prompt_type': 'simple_obo_definition',\n            'namespace_distribution': go_data['namespace'].value_counts().to_dict()\n        }\n    }\n\n    extractor.save_results(results, prefix=\"simple_go_embeddings\")\n    return results","metadata":{"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"","metadata":{"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"## Read Data and compute embeddings","metadata":{}},{"cell_type":"code","source":"obo_path = \"/kaggle/input/cafa-6-protein-function-prediction/Train/go-basic.obo\"\nresults = extract_go_embeddings(obo_path, max_terms=100)\n","metadata":{"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# simple_results = extract_simple_comparison(obo_path, max_terms=100)\n","metadata":{"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"analyze_embeddings(results)\n","metadata":{"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"print(\"\\nComparison:\")\nfor i in range(3):\n    print(f\"\\nGO: {results['go_ids'][i]}\")\n    # print(f\"Simple: {simple_results['simple_descriptions'][i][:100]}...\")\n    print(f\"Enhanced: {results['descriptions'][i][:100]}...\")","metadata":{"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"results","metadata":{"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"\n# Example usage in notebooks:\n\"\"\"\n# Copy this entire script to your notebook and run it\n\n# Simple extraction (no LLM)\nresults = extract_go_embeddings(max_terms=100)\n\n# With Gemma LLM for enhanced descriptions\n# results = extract_go_embeddings(\n#     gemma_variant=\"2b-it\",\n#     gemma_weights_dir=\"/kaggle/input/gemma/pytorch/1.1-2b-it/1/\",\n#     max_terms=100\n# )\n\n# Compare with simple approach\nsimple_results = extract_simple_comparison(obo_path, max_terms=100)\n\n# Analyze results\nanalyze_embeddings(results)\n\n# Compare descriptions\nprint(\"\\nComparison:\")\nfor i in range(3):\n    print(f\"\\nGO: {results['go_ids'][i]}\")\n    print(f\"Simple: {simple_results['simple_descriptions'][i][:100]}...\")\n    print(f\"Enhanced: {results['descriptions'][i][:100]}...\")\n\"\"\"","metadata":{"trusted":true},"outputs":[],"execution_count":null}]}