# wumaotegan/tkam_bio architecture

> Architecture diagram of the GitHub repository wumaotegan/tkam_bio, generated by GitDiagram from its file tree, README and a sample of source files (last updated 2026-09-29). It is an AI-made overview: check details against the source.

- Interactive diagram: https://gitdiagram.com/wumaotegan/tkam_bio
- Repository: https://github.com/wumaotegan/tkam_bio

## Overview

TKAM retrieves mechanistically relevant PubChem assays for HepTox and ReproTox, curates their activity data and PFAS molecular inputs, transfers bioactivity learning through a staged multi-task curriculum, and exposes endpoint-specific prediction notebooks. The README documents major workflow stages and integrations, while sampled source confirms KG construction, assay downloading, and multi-task training details; notebook internals and some integrations are not sampled, so uncertain wiring is omitted.

Main entry points: [`code/bioassay_retrieval/workflow_hepatotoxicity_bioassay_retrieval.ipynb`](https://github.com/wumaotegan/tkam_bio/blob/main/code/bioassay_retrieval/workflow_hepatotoxicity_bioassay_retrieval.ipynb) (HepTox retrieval), [`code/bioassay_retrieval/workflow_reproductive_toxicity_bioassay_retrieval.ipynb`](https://github.com/wumaotegan/tkam_bio/blob/main/code/bioassay_retrieval/workflow_reproductive_toxicity_bioassay_retrieval.ipynb) (ReproTox retrieval), [`data`](https://github.com/wumaotegan/tkam_bio/tree/main/data) (Assay and toxicity data), [`code/bioassay_retrieval/get_info_pub.py`](https://github.com/wumaotegan/tkam_bio/blob/main/code/bioassay_retrieval/get_info_pub.py) (Assay download)

## Diagram (Mermaid)

```mermaid
flowchart TD

subgraph group_retrieval["Bioassay retrieval"]
  node_knowledgebuilder["KG and embeddings"]
  node_knowledgebase[("Bioassay knowledgebase")]
  node_heptoxretrieval["HepTox retrieval"]
  node_reprotoxretrieval["ReproTox retrieval"]
  node_assaydownload["Assay download<br/>[get_info_pub.py]"]
end

subgraph group_curation["Curation and representation"]
  node_assaycuration["Assay curation"]
  node_smilesstandardization["SMILES standardization"]
  node_pfasencoder["PFAS encoder"]
end

subgraph group_training["Curriculum training"]
  node_multitasktraining["Multi-task curriculum"]
  node_curriculumtrainer["In-vivo trainer"]
  node_trainingdata["Assay and toxicity data"]
end

subgraph group_prediction["Prediction"]
  node_trainedmodels["Endpoint models"]
  node_heptoxprediction["HepTox prediction"]
  node_reprotoxprediction["ReproTox prediction"]
end

node_researcher(("Researcher"))
node_aopwiki["AOP-Wiki"]
node_pubchem["PubChem"]
node_modelservices["Model services"]
node_assaydescriptions["Assay descriptions"]

node_researcher -->|"runs workflow"| node_heptoxretrieval
node_researcher -->|"runs workflow"| node_reprotoxretrieval
node_assaydescriptions -->|"provides descriptions"| node_knowledgebuilder
node_knowledgebuilder -.->|"requests extraction and embeddings"| node_modelservices
node_knowledgebuilder -->|"builds graph"| node_knowledgebase
node_heptoxretrieval -->|"queries AOPs"| node_aopwiki
node_reprotoxretrieval -->|"queries AOPs"| node_aopwiki
node_heptoxretrieval -->|"retrieves and expands assays"| node_knowledgebase
node_reprotoxretrieval -->|"retrieves and expands assays"| node_knowledgebase
node_heptoxretrieval -.->|"reranks and reasons"| node_modelservices
node_reprotoxretrieval -.->|"reranks and reasons"| node_modelservices
node_heptoxretrieval -->|"downloads selected assays"| node_assaydownload
node_reprotoxretrieval -->|"downloads selected assays"| node_assaydownload
node_assaydownload -->|"fetches activity records"| node_pubchem
node_assaydownload -->|"supplies raw records"| node_assaycuration
node_assaycuration -->|"writes curated tasks"| node_trainingdata
node_smilesstandardization -->|"annotates PFAS inputs"| node_trainingdata
node_trainingdata -->|"supplies PFAS SMILES"| node_pfasencoder
node_pfasencoder -->|"provides adapted encoder"| node_multitasktraining
node_trainingdata -->|"supplies assay tasks"| node_multitasktraining
node_multitasktraining -->|"passes trained stages"| node_curriculumtrainer
node_trainingdata -->|"supplies in-vivo labels"| node_curriculumtrainer
node_curriculumtrainer -->|"saves models"| node_trainedmodels
node_multitasktraining -->|"produces model artifacts"| node_trainedmodels
node_researcher -->|"submits chemicals"| node_heptoxprediction
node_researcher -->|"submits chemicals"| node_reprotoxprediction
node_heptoxprediction -->|"loads endpoint model"| node_trainedmodels
node_reprotoxprediction -->|"loads endpoint model"| node_trainedmodels

click node_assaydescriptions "https://github.com/wumaotegan/tkam_bio/tree/main/data"
click node_knowledgebuilder "https://github.com/wumaotegan/tkam_bio/blob/main/code/bioassay_retrieval/generate_knowledge_graph_qwen3_base.py"
click node_knowledgebase "https://github.com/wumaotegan/tkam_bio/tree/main/data"
click node_heptoxretrieval "https://github.com/wumaotegan/tkam_bio/blob/main/code/bioassay_retrieval/workflow_hepatotoxicity_bioassay_retrieval.ipynb"
click node_reprotoxretrieval "https://github.com/wumaotegan/tkam_bio/blob/main/code/bioassay_retrieval/workflow_reproductive_toxicity_bioassay_retrieval.ipynb"
click node_assaydownload "https://github.com/wumaotegan/tkam_bio/blob/main/code/bioassay_retrieval/get_info_pub.py"
click node_assaycuration "https://github.com/wumaotegan/tkam_bio/tree/main/code/bioassay_retrieval"
click node_smilesstandardization "https://github.com/wumaotegan/tkam_bio/blob/main/code/chembert_pfas/PFAS_SMILES_standardization.ipynb"
click node_pfasencoder "https://github.com/wumaotegan/tkam_bio/blob/main/code/chembert_pfas/ChemBERTa-100M_full_finetune.ipynb"
click node_multitasktraining "https://github.com/wumaotegan/tkam_bio/blob/main/code/multi_tasks_training/multitask-pretraining-with-pruning-FT_heptox.ipynb"
click node_curriculumtrainer "https://github.com/wumaotegan/tkam_bio/blob/main/code/multi_tasks_training/invivo_train_curriculum_save_model.py"
click node_trainingdata "https://github.com/wumaotegan/tkam_bio/tree/main/data"
click node_trainedmodels "https://github.com/wumaotegan/tkam_bio/tree/main/code/predict/predict_models"
click node_heptoxprediction "https://github.com/wumaotegan/tkam_bio/blob/main/code/predict/predict_heptox.ipynb"
click node_reprotoxprediction "https://github.com/wumaotegan/tkam_bio/blob/main/code/predict/predict_reprotox.ipynb"

classDef toneNeutral fill:#f8fafc,stroke:#334155,stroke-width:1.5px,color:#0f172a
classDef toneBlue fill:#dbeafe,stroke:#2563eb,stroke-width:1.5px,color:#172554
classDef toneAmber fill:#fef3c7,stroke:#d97706,stroke-width:1.5px,color:#78350f
classDef toneMint fill:#dcfce7,stroke:#16a34a,stroke-width:1.5px,color:#14532d
classDef toneRose fill:#ffe4e6,stroke:#e11d48,stroke-width:1.5px,color:#881337
classDef toneIndigo fill:#e0e7ff,stroke:#4f46e5,stroke-width:1.5px,color:#312e81
classDef toneTeal fill:#ccfbf1,stroke:#0f766e,stroke-width:1.5px,color:#134e4a
class node_knowledgebuilder,node_knowledgebase,node_heptoxretrieval,node_reprotoxretrieval,node_assaydownload toneBlue
class node_assaycuration,node_smilesstandardization,node_pfasencoder toneAmber
class node_multitasktraining,node_curriculumtrainer,node_trainingdata toneMint
class node_trainedmodels,node_heptoxprediction,node_reprotoxprediction toneRose
class node_researcher,node_aopwiki,node_pubchem,node_modelservices toneIndigo
class node_assaydescriptions toneTeal
```

## Components

### Bioassay retrieval

- [KG and embeddings](https://github.com/wumaotegan/tkam_bio/blob/main/code/bioassay_retrieval/generate_knowledge_graph_qwen3_base.py): `code/bioassay_retrieval/generate_knowledge_graph_qwen3_base.py`
- [Bioassay knowledgebase](https://github.com/wumaotegan/tkam_bio/tree/main/data): `data`
- [HepTox retrieval](https://github.com/wumaotegan/tkam_bio/blob/main/code/bioassay_retrieval/workflow_hepatotoxicity_bioassay_retrieval.ipynb): `code/bioassay_retrieval/workflow_hepatotoxicity_bioassay_retrieval.ipynb`
- [ReproTox retrieval](https://github.com/wumaotegan/tkam_bio/blob/main/code/bioassay_retrieval/workflow_reproductive_toxicity_bioassay_retrieval.ipynb): `code/bioassay_retrieval/workflow_reproductive_toxicity_bioassay_retrieval.ipynb`
- [Assay download](https://github.com/wumaotegan/tkam_bio/blob/main/code/bioassay_retrieval/get_info_pub.py): `code/bioassay_retrieval/get_info_pub.py`

### Curation and representation

- [Assay curation](https://github.com/wumaotegan/tkam_bio/tree/main/code/bioassay_retrieval): `code/bioassay_retrieval`
- [SMILES standardization](https://github.com/wumaotegan/tkam_bio/blob/main/code/chembert_pfas/PFAS_SMILES_standardization.ipynb): `code/chembert_pfas/PFAS_SMILES_standardization.ipynb`
- [PFAS encoder](https://github.com/wumaotegan/tkam_bio/blob/main/code/chembert_pfas/ChemBERTa-100M_full_finetune.ipynb): `code/chembert_pfas/ChemBERTa-100M_full_finetune.ipynb`

### Curriculum training

- [Multi-task curriculum](https://github.com/wumaotegan/tkam_bio/blob/main/code/multi_tasks_training/multitask-pretraining-with-pruning-FT_heptox.ipynb): `code/multi_tasks_training/multitask-pretraining-with-pruning-FT_heptox.ipynb`
- [In-vivo trainer](https://github.com/wumaotegan/tkam_bio/blob/main/code/multi_tasks_training/invivo_train_curriculum_save_model.py): `code/multi_tasks_training/invivo_train_curriculum_save_model.py`
- [Assay and toxicity data](https://github.com/wumaotegan/tkam_bio/tree/main/data): `data`

### Prediction

- [Endpoint models](https://github.com/wumaotegan/tkam_bio/tree/main/code/predict/predict_models): `code/predict/predict_models`
- [HepTox prediction](https://github.com/wumaotegan/tkam_bio/blob/main/code/predict/predict_heptox.ipynb): `code/predict/predict_heptox.ipynb`
- [ReproTox prediction](https://github.com/wumaotegan/tkam_bio/blob/main/code/predict/predict_reprotox.ipynb): `code/predict/predict_reprotox.ipynb`

### Other components

- **Researcher**
- **AOP-Wiki**
- **PubChem**
- **Model services**
- [Assay descriptions](https://github.com/wumaotegan/tkam_bio/tree/main/data): `data`

## Connections

- Researcher → HepTox retrieval: runs workflow
- Researcher → ReproTox retrieval: runs workflow
- Assay descriptions → KG and embeddings: provides descriptions
- KG and embeddings → Model services: requests extraction and embeddings
- KG and embeddings → Bioassay knowledgebase: builds graph
- HepTox retrieval → AOP-Wiki: queries AOPs
- ReproTox retrieval → AOP-Wiki: queries AOPs
- HepTox retrieval → Bioassay knowledgebase: retrieves and expands assays
- ReproTox retrieval → Bioassay knowledgebase: retrieves and expands assays
- HepTox retrieval → Model services: reranks and reasons
- ReproTox retrieval → Model services: reranks and reasons
- HepTox retrieval → Assay download: downloads selected assays
- ReproTox retrieval → Assay download: downloads selected assays
- Assay download → PubChem: fetches activity records
- Assay download → Assay curation: supplies raw records
- Assay curation → Assay and toxicity data: writes curated tasks
- SMILES standardization → Assay and toxicity data: annotates PFAS inputs
- Assay and toxicity data → PFAS encoder: supplies PFAS SMILES
- PFAS encoder → Multi-task curriculum: provides adapted encoder
- Assay and toxicity data → Multi-task curriculum: supplies assay tasks
- Multi-task curriculum → In-vivo trainer: passes trained stages
- Assay and toxicity data → In-vivo trainer: supplies in-vivo labels
- In-vivo trainer → Endpoint models: saves models
- Multi-task curriculum → Endpoint models: produces model artifacts
- Researcher → HepTox prediction: submits chemicals
- Researcher → ReproTox prediction: submits chemicals
- HepTox prediction → Endpoint models: loads endpoint model
- ReproTox prediction → Endpoint models: loads endpoint model