feat: citation generator and updated pipeline

Author: shloknatarajan

data/annotations/PMC11730665.json

@@ -255,11 +255,11 @@
    ],
    "source": [
     "# Test the new PharmacogenomicTableGenerator component\n",
-    "from src.components.pharmacogenomic_table import PharmacogenomicTableGenerator\n",
+    "from src.components.annotation_table import AnnotationTableGenerator\n",
     "import json\n",
     "\n",
     "# Initialize the generator\n",
-    "table_generator = PharmacogenomicTableGenerator(pmcid=\"PMC11730665\", model=\"gpt-4.1\")\n",
+    "table_generator = AnnotationTableGenerator(pmcid=\"PMC11730665\", model=\"gpt-4.1\")\n",
     "\n",
     "# Generate structured JSON directly\n",
     "print(\"=== Generating structured JSON directly ===\")\n",
@@ -273,8 +273,54 @@
    ]
   },
   {
-   "cell_type": "markdown",
+   "cell_type": "code",
+   "execution_count": 14,
+   "metadata": {},
+   "outputs": [
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "Generating 1 Responses:   0%|          | 0/1 [00:00<?, ?it/s]\u001b[32m2025-07-29 10:09:48.996\u001b[0m | \u001b[1mINFO    \u001b[0m | \u001b[36msrc.article_parser\u001b[0m:\u001b[36m__init__\u001b[0m:\u001b[36m43\u001b[0m - \u001b[1mGetting article text from PMCID: PMC5728534\u001b[0m\n",
+      "\u001b[32m2025-07-29 10:09:48.999\u001b[0m | \u001b[1mINFO    \u001b[0m | \u001b[36msrc.article_parser\u001b[0m:\u001b[36mremove_references_section\u001b[0m:\u001b[36m90\u001b[0m - \u001b[1mRemoved References section from article text\u001b[0m\n",
+      "Generating 1 Responses: 100%|██████████| 1/1 [00:11<00:00, 11.16s/it]"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "| Gene | Polymorphism | Relationship/Effect | p-value |\n",
+      "|------|-------------|---------------------|---------|\n",
+      "| OCT1 (SLC22A1) | OCT1*5, OCT1*6, OCT1*12, OCT1*13 | These alleles completely lacked ranitidine uptake (loss of function for ranitidine transport). | Not explicitly stated, but effect is described as complete loss. |\n",
+      "| OCT1 (SLC22A1) | OCT1*2, OCT1*3, OCT1*4, OCT1*10 | These alleles had vmax values for ranitidine uptake decreased by more than 50% compared to reference. | *OCT1*2: p<0.001; *OCT1*3: p<0.001; *OCT1*4: p<0.001; *OCT1*10: p<0.05 (from Table 2) |\n",
+      "| OCT1 (SLC22A1) | OCT1*8 | This allele showed an increase of vmax for ranitidine uptake by 25% (not statistically significant). | P = 0.5 (not significant) |\n",
+      "| OCT1 (SLC22A1) | OCT1*1A, OCT1*1C, OCT1*1D, OCT1*7, OCT1*9, OCT1*11 | These alleles showed no significant difference in ranitidine uptake compared to the reference allele. | Not significant (exact values not given) |\n",
+      "| OCT1 (SLC22A1) | OCT1*2 (Met420del) | Ranitidine was on average two-fold more potent in inhibiting morphine uptake in the OCT1*2 allele compared to the reference allele (IC50 for morphine uptake: 19.5 μM for OCT1*2 vs 45.5 μM for reference). | Not explicitly stated, but difference is described as significant. |\n",
+      "| OCT2 (SLC22A2) | Ala270Ser | OCT2 showed only a limited uptake of ranitidine that was not significantly affected by the Ala270Ser polymorphism (9% reduction, not significant). | Not significant (exact value not given) |\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "\n"
+     ]
+    }
+   ],
+   "source": [
+    "from src.components.annotation_table import AnnotationTableGenerator\n",
+    "import json\n",
+    "table_gen_2 = AnnotationTableGenerator(pmcid=\"PMC5728534\", model=\"gpt-4.1\")\n",
+    "\n",
+    "table_gen_2.print_table_markdown(table_gen_2.generate_table_json())"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
    "metadata": {},
+   "outputs": [],
    "source": []
   }
  ],

data/annotations/PMC5728534.json

@@ -18,6 +18,7 @@ download-articles = "python -m src.fetch_articles.article_downloader"
 download-data = "gdown --fuzzy https://drive.google.com/file/d/1qtQWvi0x_k5_JofgrfsgkWzlIdb6isr9/view && unzip autogkb-data.zip && rm autogkb-data.zip"
 setup-repo = "pixi install && pixi run download-data"
 copy-markdown = "python -m src.copy_markdown"
+annotation-pipeline = "python -m src.components.annotation_pipeline"
 
 [dependencies]
 seaborn = ">=0.13.2,<0.14"

notebooks/table_generation.ipynb

@@ -1,61 +1,49 @@
-from src.components.all_associations import get_all_associations, AssociationType
-from src.components.drug_annotation import get_drug_annotation
-from src.components.phenotype_annotation import get_phenotype_annotation
-from src.components.functional_annotation import get_functional_annotation
+from src.components.annotation_table import AnnotationTableGenerator
+from src.components.citation_generator import CitationGenerator
 from src.components.study_parameters import get_study_parameters
 from src.utils import get_article_text, is_pmcid, get_title
-from typing import Optional
 from loguru import logger
 from pathlib import Path
 import os
 
 class AnnotationPipeline:
-    def __init__(self, pmcid: str):
+    def __init__(self, pmcid: str, citation_approach: str = "lm"):
         if not is_pmcid(pmcid):
             logger.error(f"Invalid PMCID: {pmcid}")
         self.pmcid = pmcid
+        self.citation_approach = citation_approach
         self.article_text = get_article_text(pmcid)
         self.title = get_title(self.article_text)
-        self.all_associations = []
         self.study_parameters = {}
-        self.drug_annotations = []
-        self.phenotye_annotations = []
-        self.functional_annotations = []
+        self.annotations = None
 
     def print_info(self):
-        logger.info(f"Found {len(self.all_associations)} associations")
-        logger.info(f"Created {len(self.drug_annotations)} Drug Annotations")
-        logger.info(f"Created {len(self.phenotye_annotations)} Phenotype Annotations")
-        logger.info(
-            f"Created {len(self.functional_annotations)} Functional Annotations"
-        )
+        annotation_count = len(self.annotations.relationships) if self.annotations else 0
+        
+        logger.info(f"Created {annotation_count} Annotations")
 
     def generate_final_structure(self):
         return {
             "pmcid": self.pmcid,
             "title": self.title,
             "study_parameters": self.study_parameters,
-            "drug_annotations": self.drug_annotations,
-            "phenotype_annotations": self.phenotye_annotations,
-            "functional_annotations": self.functional_annotations,
+            "annotations": self.annotations,
         }
 
     def run(self, save_path: str = "data/annotations"):
         logger.info("Getting Study Parameters")
         self.study_parameters = get_study_parameters(self.article_text)
 
-        logger.info("Getting All Associations")
-        self.all_associations = get_all_associations(self.article_text)
+        # Generate annotations using AnnotationTableGenerator
+        annotation_generator = AnnotationTableGenerator(self.pmcid)
+        
+        logger.info("Generating Annotations")
+        self.annotations = annotation_generator.generate_table_json()
 
-        for association in self.all_associations:
-            if association.association_type == AssociationType.DRUG:
-                self.drug_annotations.append(get_drug_annotation(association))
-            if association.association_type == AssociationType.PHENOTYPE:
-                self.phenotye_annotations.append(get_phenotype_annotation(association))
-            if association.association_type == AssociationType.FUNCTIONAL:
-                self.functional_annotations.append(
-                    get_functional_annotation(association)
-                )
+        # Generate citations for annotations
+        citation_generator = CitationGenerator(self.pmcid, approach=self.citation_approach)
+        logger.info(f"Adding Citations to Annotations using {self.citation_approach} approach")
+        self.annotations = citation_generator.add_citations_to_annotations(self.annotations)
 
         self.print_info()
 
@@ -89,14 +77,14 @@ def copy_markdown(pmcid: str):
 if __name__ == "__main__":
     pmcids = [
         "PMC5728534",
-        # "PMC11730665",
-        # "PMC5712579",
-        # "PMC4737107",
-        # "PMC5749368"
+        "PMC11730665",
+        "PMC5712579",
+        "PMC4737107",
+        "PMC5749368"
     ]
     for pmcid in pmcids:
         logger.info(f"Processing {pmcid}")
-        pipeline = AnnotationPipeline(pmcid)
+        pipeline = AnnotationPipeline(pmcid, citation_approach="local")
         pipeline.run()
-    # for pmcid in pmcids:
-    #     copy_markdown(pmcid)
+    for pmcid in pmcids:
+        copy_markdown(pmcid)

pixi.toml

@@ -9,6 +9,7 @@ class AnnotationRelationship(BaseModel):
     polymorphism: str = Field(description="Genetic polymorphism/variant")
     relationship_effect: str = Field(description="Relationship or effect description")
     p_value: str = Field(description="Statistical p-value")
+    citations: List[str] = Field(default=[], description="List of supporting sentences from the text")
 
 
 class AnnotationTable(BaseModel):