Add files using upload-large-folder tool

.gitattributes

@@ -39,3 +39,18 @@ all_checkpoints/stage2_07070337_2datasets_noconstruct/wandb/run-20250707_041231-
 all_checkpoints/stage2_07070513_2datasets_construct/wandb/run-20250707_052104-615z4bme/run-615z4bme.wandb filter=lfs diff=lfs merge=lfs -text
 all_checkpoints/stage2_07070513_2datasets_construct/wandb/run-20250707_053222-9cjzn0v3/run-9cjzn0v3.wandb filter=lfs diff=lfs merge=lfs -text
 all_checkpoints/stage2_07301646_2datasets_construct/wandb/run-20250730_175623-pbf2bxo6/run-pbf2bxo6.wandb filter=lfs diff=lfs merge=lfs -text
+data/OntoProteinDatasetV2/test.txt filter=lfs diff=lfs merge=lfs -text
+all_checkpoints/stage2_08011616_2datasets_withoutpretrain/wandb/run-20250801_162632-790wkw0g/run-790wkw0g.wandb filter=lfs diff=lfs merge=lfs -text
+data/OntoProteinDatasetV2/train.txt filter=lfs diff=lfs merge=lfs -text
+data/OntoProteinDatasetV2/valid.txt filter=lfs diff=lfs merge=lfs -text
+data/PDBDataset/abstract.json filter=lfs diff=lfs merge=lfs -text
+data/PDBDataset/qa_all.json filter=lfs diff=lfs merge=lfs -text
+data/PDBDataset/train.txt filter=lfs diff=lfs merge=lfs -text
+data/SwissProtV3/train_set.jsonl filter=lfs diff=lfs merge=lfs -text
+data/protein-text/eval_assist.ziphwjr8q2y.tmp filter=lfs diff=lfs merge=lfs -text
+data/protein-text/eval_assist.zipg3ebgjl7.tmp filter=lfs diff=lfs merge=lfs -text
+data/protein-text/eval_assist.zipzh1pdmj_.tmp filter=lfs diff=lfs merge=lfs -text
+data_small/OntoProteinDatasetV2/train.txt filter=lfs diff=lfs merge=lfs -text
+data_small/PDBDataset/abstract.json filter=lfs diff=lfs merge=lfs -text
+data_small/PDBDataset/qa_all.json filter=lfs diff=lfs merge=lfs -text
+data_small/SwissProtV3/train_set_.jsonl filter=lfs diff=lfs merge=lfs -text

all_checkpoints/stage2_08011616_2datasets_withoutpretrain/last.ckpt/checkpoint/mp_rank_00_model_states.pt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:6efbc882065731ceb2c9886091da92484e90352ef036f1fff44e77700ff80f41
+size 208795384

all_checkpoints/stage2_08011616_2datasets_withoutpretrain/wandb/run-20250801_162632-790wkw0g/files/requirements.txt

@@ -0,0 +1,225 @@
+opendatasets==0.1.22
+salesforce-lavis==1.0.2
+Pygments==2.19.1
+nvidia-nccl-cu12==2.21.5
+tornado==6.5.1
+nvidia-cuda-runtime-cu12==12.4.127
+requests==2.32.3
+nvidia-cuda-cupti-cu12==12.4.127
+decord==0.6.0
+braceexpand==0.1.7
+frozenlist==1.6.0
+markdown-it-py==3.0.0
+shellingham==1.5.4
+absl-py==2.2.2
+pycocoevalcap==1.2
+contexttimer==0.3.3
+bleach==6.2.0
+jsonschema-specifications==2025.4.1
+pycocotools==2.0.8
+python-slugify==8.0.4
+tqdm==4.67.1
+numpy==2.2.6
+urllib3==2.4.0
+deepspeed==0.16.10+b666844f
+watchdog==6.0.0
+wrapt==1.17.2
+setuptools==78.1.1
+matplotlib==3.10.3
+pydeck==0.9.1
+aiosignal==1.3.2
+gitdb==4.0.12
+hjson==3.1.0
+timm==0.4.12
+blis==1.3.0
+PyYAML==6.0.2
+referencing==0.36.2
+contourpy==1.3.2
+kaggle==1.7.4.5
+triton==3.2.0
+catalogue==2.0.10
+idna==3.10
+torch==2.6.0
+text-unidecode==1.3
+altair==5.5.0
+cloudpathlib==0.21.1
+protobuf==6.31.0
+nvidia-cusolver-cu12==11.6.1.9
+pytz==2025.2
+sympy==1.13.1
+spacy==3.8.7
+MarkupSafe==3.0.2
+thinc==8.3.6
+nvidia-cudnn-cu12==9.1.0.70
+wasabi==1.1.3
+aiohappyeyeballs==2.6.1
+nvidia-nvtx-cu12==12.4.127
+rich==14.0.0
+ipython==8.36.0
+yarl==1.20.0
+torchmetrics==1.7.1
+multidict==6.4.4
+cfgv==3.4.0
+smmap==5.0.2
+srsly==2.5.1
+scikit-image==0.25.2
+matplotlib-inline==0.1.7
+annotated-types==0.7.0
+lazy_loader==0.4
+tenacity==9.1.2
+GitPython==3.1.44
+language_data==1.3.0
+pydantic_core==2.33.2
+sentencepiece==0.2.0
+platformdirs==4.3.8
+distlib==0.3.9
+nvidia-cusparselt-cu12==0.6.2
+blinker==1.9.0
+regex==2024.11.6
+tifffile==2025.5.10
+py-cpuinfo==9.0.0
+attrs==25.3.0
+mdurl==0.1.2
+prompt_toolkit==3.0.51
+packaging==24.2
+async-timeout==5.0.1
+six==1.17.0
+executing==2.2.0
+parso==0.8.4
+omegaconf==2.3.0
+wcwidth==0.2.13
+murmurhash==1.0.13
+stack-data==0.6.3
+nvidia-cufft-cu12==11.2.1.3
+virtualenv==20.31.2
+langcodes==3.5.0
+fonttools==4.58.0
+opencv-python-headless==4.5.5.64
+jedi==0.19.2
+torchvision==0.21.0
+plotly==6.1.1
+nodeenv==1.9.1
+smart-open==7.1.0
+toml==0.10.2
+pytorch-lightning==2.5.1.post0
+typing_extensions==4.13.2
+safetensors==0.5.3
+psutil==7.0.0
+pillow==11.2.1
+python-dateutil==2.9.0.post0
+ftfy==6.3.1
+scipy==1.15.3
+webdataset==0.2.111
+charset-normalizer==3.4.2
+nvidia-nvjitlink-cu12==12.4.127
+kiwisolver==1.4.8
+nvidia-ml-py==12.575.51
+confection==0.1.5
+nvidia-curand-cu12==10.3.5.147
+pandas==2.2.3
+nltk==3.9.1
+webencodings==0.5.1
+pyarrow==20.0.0
+asttokens==3.0.0
+exceptiongroup==1.3.0
+pre_commit==4.2.0
+ninja==1.11.1.4
+spacy-loggers==1.0.5
+msgpack==1.1.0
+lightning-utilities==0.14.3
+nvidia-cublas-cu12==12.4.5.8
+tzdata==2025.2
+cycler==0.12.1
+hf-xet==1.1.2
+antlr4-python3-runtime==4.9.3
+iopath==0.1.10
+pexpect==4.9.0
+imageio==2.37.0
+streamlit==1.45.1
+python-magic==0.4.27
+networkx==3.4.2
+portalocker==3.1.1
+nvidia-cusparse-cu12==12.3.1.170
+propcache==0.3.1
+ptyprocess==0.7.0
+fairscale==0.4.4
+rpds-py==0.25.1
+certifi==2025.4.26
+rouge_score==0.1.2
+traitlets==5.14.3
+identify==2.6.12
+spacy-legacy==3.0.12
+weasel==0.4.1
+mpmath==1.3.0
+cymem==2.0.11
+typing-inspection==0.4.1
+nvidia-cuda-nvrtc-cu12==12.4.127
+marisa-trie==1.2.1
+einops==0.8.1
+nvidia-cufile-cu12==1.11.1.6
+pydantic==2.11.5
+cachetools==5.5.2
+joblib==1.5.1
+Jinja2==3.1.6
+filelock==3.18.0
+pyparsing==3.2.3
+pure_eval==0.2.3
+decorator==5.2.1
+wheel==0.45.1
+pycryptodome==3.23.0
+cheroot==10.0.1
+multiprocess==0.70.16
+aiohttp==3.12.2
+crcmod==1.7
+fsspec==2025.3.0
+jmespath==0.10.0
+preshed==3.0.10
+jaraco.functools==4.1.0
+cryptography==45.0.3
+sentry-sdk==2.29.1
+tokenizers==0.21.1
+opendelta==0.3.2
+pycparser==2.22
+narwhals==1.41.0
+scikit-learn==1.6.1
+dill==0.3.8
+oss2==2.15.0
+yacs==0.1.8
+more-itertools==10.7.0
+pip==25.1.1
+threadpoolctl==3.6.0
+flash-attn==2.7.1.post1
+bigmodelvis==0.0.1
+pathlib==1.0.1
+delta-center-client==0.0.4
+xxhash==3.5.0
+wandb==0.19.11
+setproctitle==1.3.6
+aliyun-python-sdk-core==2.16.0
+transformers==4.52.3
+aliyun-python-sdk-kms==2.16.5
+datasets==3.6.0
+typer==0.16.0
+docker-pycreds==0.4.0
+click==8.2.1
+huggingface-hub==0.32.1
+web.py==0.62
+cffi==1.17.1
+opencv-python==4.11.0.86
+jsonschema==4.24.0
+typing_extensions==4.12.2
+jaraco.functools==4.0.1
+jaraco.text==3.12.1
+jaraco.collections==5.1.0
+inflect==7.3.1
+more-itertools==10.3.0
+packaging==24.2
+importlib_metadata==8.0.0
+backports.tarfile==1.2.0
+typeguard==4.3.0
+zipp==3.19.2
+platformdirs==4.2.2
+autocommand==2.2.2
+jaraco.context==5.3.0
+tomli==2.0.1
+wheel==0.45.1

all_checkpoints/stage2_08011616_2datasets_withoutpretrain/wandb/run-20250801_162632-790wkw0g/files/wandb-metadata.json

@@ -0,0 +1,108 @@
+{
+  "os": "Linux-5.10.134-008.16.kangaroo.al8.x86_64-x86_64-with-glibc2.35",
+  "python": "CPython 3.10.0",
+  "startedAt": "2025-08-01T08:26:32.241935Z",
+  "args": [
+    "--devices",
+    "0,1,2,3,4,5,6,7",
+    "--mode",
+    "train",
+    "--filename",
+    "stage2_08011616_2datasets_qweninstruct",
+    "--num_query_token",
+    "8",
+    "--save_every_n_epochs",
+    "2",
+    "--max_epochs",
+    "10",
+    "--batch_size",
+    "4",
+    "--precision",
+    "bf16-mixed",
+    "--num_workers",
+    "8",
+    "--plm_model",
+    "/nas/shared/kilab/wangyujia/ProtT3/plm_model/esm2-150m",
+    "--bert_name",
+    "/nas/shared/kilab/wangyujia/ProtT3/plm_model/microsoft",
+    "--llm_name",
+    "/nas/shared/kilab/hf-hub/models--Qwen--Qwen2.5-7B-Instruct/snapshots/a09a35458c702b33eeacc393d103063234e8bc28",
+    "--llm_tune",
+    "mid_lora",
+    "--stage1_path",
+    "/nas/shared/kilab/wangyujia/ProtT3/all_checkpoints/stage1_07041727_2dataset/epoch=29.ckpt/converted.ckpt",
+    "--use_wandb_logger",
+    "--dataset",
+    "swiss-prot"
+  ],
+  "program": "/nas/shared/kilab/wangyujia/ProtT3/stage2.py",
+  "codePath": "wangyujia/ProtT3/stage2.py",
+  "git": {
+    "remote": "https://github.com/PorUna-byte/PAR.git",
+    "commit": "b8caf406aa1699c788f0ca6e44a1769452c317db"
+  },
+  "root": "./all_checkpoints/stage2_08011616_2datasets_qweninstruct/",
+  "host": "dsw-265304-58fbcf9d9b-zvtdx",
+  "executable": "/root/miniconda3/envs/protT3/bin/python",
+  "codePathLocal": "stage2.py",
+  "cpu_count": 64,
+  "cpu_count_logical": 64,
+  "gpu": "NVIDIA A800-SXM4-80GB",
+  "gpu_count": 8,
+  "disk": {
+    "/": {
+      "total": "1623302262784",
+      "used": "1476915200"
+    }
+  },
+  "memory": {
+    "total": "549755813888"
+  },
+  "cpu": {
+    "count": 64,
+    "countLogical": 64
+  },
+  "gpu_nvidia": [
+    {
+      "name": "NVIDIA A800-SXM4-80GB",
+      "memoryTotal": "85198045184",
+      "architecture": "Ampere"
+    },
+    {
+      "name": "NVIDIA A800-SXM4-80GB",
+      "memoryTotal": "85198045184",
+      "architecture": "Ampere"
+    },
+    {
+      "name": "NVIDIA A800-SXM4-80GB",
+      "memoryTotal": "85198045184",
+      "architecture": "Ampere"
+    },
+    {
+      "name": "NVIDIA A800-SXM4-80GB",
+      "memoryTotal": "85198045184",
+      "architecture": "Ampere"
+    },
+    {
+      "name": "NVIDIA A800-SXM4-80GB",
+      "memoryTotal": "85198045184",
+      "architecture": "Ampere"
+    },
+    {
+      "name": "NVIDIA A800-SXM4-80GB",
+      "memoryTotal": "85198045184",
+      "architecture": "Ampere"
+    },
+    {
+      "name": "NVIDIA A800-SXM4-80GB",
+      "memoryTotal": "85198045184",
+      "architecture": "Ampere"
+    },
+    {
+      "name": "NVIDIA A800-SXM4-80GB",
+      "memoryTotal": "85198045184",
+      "architecture": "Ampere"
+    }
+  ],
+  "cudaVersion": "12.1"
+}

all_checkpoints/stage2_08011616_2datasets_withoutpretrain/wandb/run-20250801_162632-790wkw0g/files/wandb-summary.json

@@ -0,0 +1 @@
+{"trainer/global_step":134559,"dataset0/rouge_2":37.98429870605469,"_step":2700,"dataset0/bleu2":37.878440856933594,"dataset0/meteor_score":53.153656005859375,"lr":1.2202456673549023e-05,"dataset0/rouge_1":45.05393981933594,"_wandb":{"runtime":61483},"_timestamp":1.7540982570645058e+09,"dataset0/rouge_l":42.82152557373047,"dataset0/bleu4":34.96389389038086,"_runtime":61467.758475346,"epoch":9,"loss":0.37739327549934387,"dataloader0/val loss/dataloader_idx_0":0.6249951124191284,"dataset0/acc":0}

all_checkpoints/stage2_08011616_2datasets_withoutpretrain/wandb/run-20250801_162632-790wkw0g/logs/debug-internal.log

@@ -0,0 +1,166 @@
+{"time":"2025-08-01T16:26:32.258940024+08:00","level":"INFO","msg":"stream: starting","core version":"0.19.11","symlink path":"all_checkpoints/stage2_08011616_2datasets_qweninstruct/wandb/run-20250801_162632-790wkw0g/logs/debug-core.log"}
+{"time":"2025-08-01T16:26:47.100144119+08:00","level":"INFO","msg":"created new stream","id":"790wkw0g"}
+{"time":"2025-08-01T16:26:47.101616008+08:00","level":"INFO","msg":"stream: started","id":"790wkw0g"}
+{"time":"2025-08-01T16:26:47.101680831+08:00","level":"INFO","msg":"sender: started","stream_id":"790wkw0g"}
+{"time":"2025-08-01T16:26:47.101644151+08:00","level":"INFO","msg":"writer: Do: started","stream_id":"790wkw0g"}
+{"time":"2025-08-01T16:26:47.101748992+08:00","level":"INFO","msg":"handler: started","stream_id":"790wkw0g"}
+{"time":"2025-08-01T16:26:50.002107204+08:00","level":"INFO","msg":"Starting system monitor"}
+{"time":"2025-08-01T16:31:09.07203296+08:00","level":"INFO","msg":"api: retrying error","error":"Post \"https://api.bandw.top/files/gia0603yucca/stage2_08011616_2datasets_qweninstruct/790wkw0g/file_stream\": read tcp 10.1.11.100:36310->104.21.20.172:443: read: connection timed out"}
+{"time":"2025-08-01T16:38:20.137992503+08:00","level":"INFO","msg":"api: retrying error","error":"Post \"https://api.bandw.top/graphql\": net/http: request canceled (Client.Timeout exceeded while awaiting headers)"}
+{"time":"2025-08-01T16:41:40.88100016+08:00","level":"INFO","msg":"api: retrying error","error":"Post \"https://api.bandw.top/files/gia0603yucca/stage2_08011616_2datasets_qweninstruct/790wkw0g/file_stream\": read tcp 10.1.11.100:41026->104.21.20.172:443: read: connection timed out"}
+{"time":"2025-08-01T16:43:35.141344346+08:00","level":"INFO","msg":"api: retrying error","error":"Post \"https://api.bandw.top/graphql\": net/http: request canceled (Client.Timeout exceeded while awaiting headers)"}
+{"time":"2025-08-01T16:44:07.544267418+08:00","level":"INFO","msg":"api: retrying error","error":"Post \"https://api.bandw.top/graphql\": net/http: request canceled (Client.Timeout exceeded while awaiting headers)"}
+{"time":"2025-08-01T16:44:41.670680082+08:00","level":"INFO","msg":"api: retrying error","error":"Post \"https://api.bandw.top/graphql\": context deadline exceeded"}
+{"time":"2025-08-01T16:45:20.186411207+08:00","level":"INFO","msg":"api: retrying error","error":"Post \"https://api.bandw.top/graphql\": context deadline exceeded"}
+{"time":"2025-08-01T16:45:57.904074825+08:00","level":"INFO","msg":"api: retrying error","error":"Post \"https://api.bandw.top/files/gia0603yucca/stage2_08011616_2datasets_qweninstruct/790wkw0g/file_stream\": read tcp 10.1.11.100:57230->172.67.193.61:443: read: connection timed out"}
+{"time":"2025-08-01T16:46:09.248932426+08:00","level":"INFO","msg":"api: retrying error","error":"Post \"https://api.bandw.top/graphql\": net/http: request canceled (Client.Timeout exceeded while awaiting headers)"}
+{"time":"2025-08-01T16:50:22.097023685+08:00","level":"INFO","msg":"api: retrying error","error":"Post \"https://api.bandw.top/files/gia0603yucca/stage2_08011616_2datasets_qweninstruct/790wkw0g/file_stream\": read tcp 10.1.11.100:45466->104.21.20.172:443: read: connection timed out"}
+{"time":"2025-08-01T17:44:15.888009977+08:00","level":"INFO","msg":"api: retrying error","error":"Post \"https://api.bandw.top/files/gia0603yucca/stage2_08011616_2datasets_qweninstruct/790wkw0g/file_stream\": read tcp 10.1.11.100:47510->104.21.20.172:443: read: connection timed out"}
+{"time":"2025-08-01T20:05:52.724089483+08:00","level":"INFO","msg":"api: retrying error","error":"Post \"https://api.bandw.top/files/gia0603yucca/stage2_08011616_2datasets_qweninstruct/790wkw0g/file_stream\": read tcp 10.1.11.100:49218->104.21.20.172:443: read: connection reset by peer"}
+{"time":"2025-08-01T20:13:58.928063914+08:00","level":"INFO","msg":"api: retrying error","error":"Post \"https://api.bandw.top/files/gia0603yucca/stage2_08011616_2datasets_qweninstruct/790wkw0g/file_stream\": read tcp 10.1.11.100:46332->104.21.20.172:443: read: connection timed out"}
+{"time":"2025-08-01T20:27:15.089020249+08:00","level":"INFO","msg":"api: retrying error","error":"Post \"https://api.bandw.top/files/gia0603yucca/stage2_08011616_2datasets_qweninstruct/790wkw0g/file_stream\": read tcp 10.1.11.100:51288->104.21.20.172:443: read: connection timed out"}
+{"time":"2025-08-01T20:33:59.569042011+08:00","level":"INFO","msg":"api: retrying error","error":"Post \"https://api.bandw.top/files/gia0603yucca/stage2_08011616_2datasets_qweninstruct/790wkw0g/file_stream\": read tcp 10.1.11.100:38854->172.67.193.61:443: read: connection timed out"}
+{"time":"2025-08-01T20:34:32.461391926+08:00","level":"INFO","msg":"api: retrying error","error":"Post \"https://api.bandw.top/files/gia0603yucca/stage2_08011616_2datasets_qweninstruct/790wkw0g/file_stream\": unexpected EOF"}
+{"time":"2025-08-01T20:38:30.928056966+08:00","level":"INFO","msg":"api: retrying error","error":"Post \"https://api.bandw.top/files/gia0603yucca/stage2_08011616_2datasets_qweninstruct/790wkw0g/file_stream\": read tcp 10.1.11.100:45810->104.21.20.172:443: read: connection timed out"}
+{"time":"2025-08-01T20:43:03.824002354+08:00","level":"INFO","msg":"api: retrying error","error":"Post \"https://api.bandw.top/files/gia0603yucca/stage2_08011616_2datasets_qweninstruct/790wkw0g/file_stream\": read tcp 10.1.11.100:51984->104.21.20.172:443: read: connection timed out"}
+{"time":"2025-08-01T20:46:25.450089946+08:00","level":"INFO","msg":"api: retrying error","error":"Post \"https://api.bandw.top/files/gia0603yucca/stage2_08011616_2datasets_qweninstruct/790wkw0g/file_stream\": read tcp 10.1.11.100:42776->172.67.193.61:443: read: connection reset by peer"}
+{"time":"2025-08-01T20:49:05.052996031+08:00","level":"INFO","msg":"api: retrying error","error":"Post \"https://api.bandw.top/files/gia0603yucca/stage2_08011616_2datasets_qweninstruct/790wkw0g/file_stream\": read tcp 10.1.11.100:54072->104.21.20.172:443: read: connection reset by peer"}
+{"time":"2025-08-01T20:49:50.926181978+08:00","level":"INFO","msg":"api: retrying error","error":"Post \"https://api.bandw.top/graphql\": net/http: request canceled (Client.Timeout exceeded while awaiting headers)"}
+{"time":"2025-08-01T20:50:58.961126603+08:00","level":"INFO","msg":"api: retrying error","error":"Post \"https://api.bandw.top/files/gia0603yucca/stage2_08011616_2datasets_qweninstruct/790wkw0g/file_stream\": read tcp 10.1.11.100:46418->104.21.20.172:443: read: connection reset by peer"}
+{"time":"2025-08-01T20:54:00.208036179+08:00","level":"INFO","msg":"api: retrying error","error":"Post \"https://api.bandw.top/files/gia0603yucca/stage2_08011616_2datasets_qweninstruct/790wkw0g/file_stream\": read tcp 10.1.11.100:47556->104.21.20.172:443: read: connection timed out"}
+{"time":"2025-08-01T21:03:33.648034502+08:00","level":"INFO","msg":"api: retrying error","error":"Post \"https://api.bandw.top/files/gia0603yucca/stage2_08011616_2datasets_qweninstruct/790wkw0g/file_stream\": read tcp 10.1.11.100:45288->104.21.20.172:443: read: connection timed out"}
+{"time":"2025-08-01T21:06:14.416001149+08:00","level":"INFO","msg":"api: retrying error","error":"Post \"https://api.bandw.top/files/gia0603yucca/stage2_08011616_2datasets_qweninstruct/790wkw0g/file_stream\": read tcp 10.1.11.100:39640->172.67.193.61:443: read: connection timed out"}
+{"time":"2025-08-01T21:06:50.935397642+08:00","level":"INFO","msg":"api: retrying error","error":"Post \"https://api.bandw.top/graphql\": context deadline exceeded"}
+{"time":"2025-08-01T21:07:22.991031701+08:00","level":"INFO","msg":"api: retrying error","error":"Post \"https://api.bandw.top/graphql\": net/http: request canceled (Client.Timeout exceeded while awaiting headers)"}
+{"time":"2025-08-01T21:07:57.40524595+08:00","level":"INFO","msg":"api: retrying error","error":"Post \"https://api.bandw.top/graphql\": context deadline exceeded"}
+{"time":"2025-08-01T21:08:36.813512769+08:00","level":"INFO","msg":"api: retrying error","error":"Post \"https://api.bandw.top/graphql\": net/http: request canceled (Client.Timeout exceeded while awaiting headers)"}
+{"time":"2025-08-01T21:09:24.678708826+08:00","level":"INFO","msg":"api: retrying error","error":"Post \"https://api.bandw.top/graphql\": net/http: request canceled (Client.Timeout exceeded while awaiting headers)"}
+{"time":"2025-08-01T21:10:33.369865787+08:00","level":"INFO","msg":"api: retrying error","error":"Post \"https://api.bandw.top/graphql\": net/http: request canceled (Client.Timeout exceeded while awaiting headers)"}
+{"time":"2025-08-01T21:13:09.331256642+08:00","level":"INFO","msg":"api: retrying error","error":"Post \"https://api.bandw.top/files/gia0603yucca/stage2_08011616_2datasets_qweninstruct/790wkw0g/file_stream\": read tcp 10.1.11.100:55554->104.21.20.172:443: read: connection reset by peer"}
+{"time":"2025-08-01T21:18:12.321384192+08:00","level":"INFO","msg":"api: retrying error","error":"Post \"https://api.bandw.top/files/gia0603yucca/stage2_08011616_2datasets_qweninstruct/790wkw0g/file_stream\": unexpected EOF"}
+{"time":"2025-08-01T21:21:14.379467148+08:00","level":"INFO","msg":"api: retrying error","error":"Post \"https://api.bandw.top/files/gia0603yucca/stage2_08011616_2datasets_qweninstruct/790wkw0g/file_stream\": net/http: request canceled (Client.Timeout exceeded while awaiting headers)"}
+{"time":"2025-08-01T21:24:32.868908249+08:00","level":"INFO","msg":"api: retrying error","error":"Post \"https://api.bandw.top/files/gia0603yucca/stage2_08011616_2datasets_qweninstruct/790wkw0g/file_stream\": context deadline exceeded"}
+{"time":"2025-08-01T21:28:13.328011263+08:00","level":"INFO","msg":"api: retrying error","error":"Post \"https://api.bandw.top/files/gia0603yucca/stage2_08011616_2datasets_qweninstruct/790wkw0g/file_stream\": read tcp 10.1.11.100:53294->104.21.20.172:443: read: connection timed out"}
+{"time":"2025-08-01T21:30:01.911292424+08:00","level":"INFO","msg":"api: retrying error","error":"Post \"https://api.bandw.top/files/gia0603yucca/stage2_08011616_2datasets_qweninstruct/790wkw0g/file_stream\": read tcp 10.1.11.100:45656->172.67.193.61:443: read: connection reset by peer"}
+{"time":"2025-08-01T21:31:05.951614465+08:00","level":"INFO","msg":"api: retrying error","error":"Post \"https://api.bandw.top/graphql\": net/http: request canceled (Client.Timeout exceeded while awaiting headers)"}
+{"time":"2025-08-01T21:32:43.664038993+08:00","level":"INFO","msg":"api: retrying error","error":"Post \"https://api.bandw.top/files/gia0603yucca/stage2_08011616_2datasets_qweninstruct/790wkw0g/file_stream\": read tcp 10.1.11.100:58998->172.67.193.61:443: read: connection timed out"}
+{"time":"2025-08-01T21:33:35.953483769+08:00","level":"ERROR","msg":"sender: sendStopStatus: failed to get run stopped status: context deadline exceeded (Client.Timeout or context cancellation while reading body)"}
+{"time":"2025-08-01T21:35:00.100184499+08:00","level":"INFO","msg":"api: retrying error","error":"Post \"https://api.bandw.top/files/gia0603yucca/stage2_08011616_2datasets_qweninstruct/790wkw0g/file_stream\": read tcp 10.1.11.100:48498->104.21.20.172:443: read: connection reset by peer"}
+{"time":"2025-08-01T21:37:35.955239872+08:00","level":"ERROR","msg":"sender: sendStopStatus: failed to get run stopped status: context deadline exceeded (Client.Timeout or context cancellation while reading body)"}
+{"time":"2025-08-01T21:38:05.712022165+08:00","level":"INFO","msg":"api: retrying error","error":"Post \"https://api.bandw.top/files/gia0603yucca/stage2_08011616_2datasets_qweninstruct/790wkw0g/file_stream\": read tcp 10.1.11.100:60866->104.21.20.172:443: read: connection timed out"}
+{"time":"2025-08-01T21:38:20.956070921+08:00","level":"INFO","msg":"api: retrying error","error":"Post \"https://api.bandw.top/graphql\": context deadline exceeded"}
+{"time":"2025-08-01T21:41:06.960060464+08:00","level":"INFO","msg":"api: retrying error","error":"Post \"https://api.bandw.top/files/gia0603yucca/stage2_08011616_2datasets_qweninstruct/790wkw0g/file_stream\": read tcp 10.1.11.100:55500->172.67.193.61:443: read: connection timed out"}
+{"time":"2025-08-01T21:41:57.697000345+08:00","level":"INFO","msg":"api: retrying error","error":"Post \"https://api.bandw.top/files/gia0603yucca/stage2_08011616_2datasets_qweninstruct/790wkw0g/file_stream\": read tcp 10.1.11.100:36556->172.67.193.61:443: read: connection reset by peer"}
+{"time":"2025-08-01T21:45:01.968042256+08:00","level":"INFO","msg":"api: retrying error","error":"Post \"https://api.bandw.top/files/gia0603yucca/stage2_08011616_2datasets_qweninstruct/790wkw0g/file_stream\": read tcp 10.1.11.100:53014->172.67.193.61:443: read: connection timed out"}
+{"time":"2025-08-01T21:46:05.960063741+08:00","level":"INFO","msg":"api: retrying error","error":"Post \"https://api.bandw.top/graphql\": net/http: request canceled (Client.Timeout exceeded while awaiting headers)"}
+{"time":"2025-08-01T21:47:02.576818619+08:00","level":"INFO","msg":"api: retrying error","error":"Post \"https://api.bandw.top/files/gia0603yucca/stage2_08011616_2datasets_qweninstruct/790wkw0g/file_stream\": read tcp 10.1.11.100:44354->104.21.20.172:443: read: connection reset by peer"}
+{"time":"2025-08-01T21:54:27.728015687+08:00","level":"INFO","msg":"api: retrying error","error":"Post \"https://api.bandw.top/files/gia0603yucca/stage2_08011616_2datasets_qweninstruct/790wkw0g/file_stream\": read tcp 10.1.11.100:59618->104.21.20.172:443: read: connection timed out"}
+{"time":"2025-08-01T21:57:21.808006007+08:00","level":"INFO","msg":"api: retrying error","error":"Post \"https://api.bandw.top/files/gia0603yucca/stage2_08011616_2datasets_qweninstruct/790wkw0g/file_stream\": read tcp 10.1.11.100:46320->104.21.20.172:443: read: connection timed out"}
+{"time":"2025-08-01T22:00:19.984010382+08:00","level":"INFO","msg":"api: retrying error","error":"Post \"https://api.bandw.top/files/gia0603yucca/stage2_08011616_2datasets_qweninstruct/790wkw0g/file_stream\": read tcp 10.1.11.100:53128->104.21.20.172:443: read: connection timed out"}
+{"time":"2025-08-01T22:00:50.968803427+08:00","level":"INFO","msg":"api: retrying error","error":"Post \"https://api.bandw.top/graphql\": net/http: request canceled (Client.Timeout exceeded while awaiting headers)"}
+{"time":"2025-08-01T22:03:12.527983291+08:00","level":"INFO","msg":"api: retrying error","error":"Post \"https://api.bandw.top/files/gia0603yucca/stage2_08011616_2datasets_qweninstruct/790wkw0g/file_stream\": read tcp 10.1.11.100:34872->104.21.20.172:443: read: connection timed out"}
+{"time":"2025-08-01T22:04:48.400864668+08:00","level":"INFO","msg":"api: retrying error","error":"Post \"https://api.bandw.top/files/gia0603yucca/stage2_08011616_2datasets_qweninstruct/790wkw0g/file_stream\": read tcp 10.1.11.100:43930->104.21.20.172:443: read: connection reset by peer"}
+{"time":"2025-08-01T22:04:50.970855218+08:00","level":"INFO","msg":"api: retrying error","error":"Post \"https://api.bandw.top/graphql\": net/http: request canceled (Client.Timeout exceeded while awaiting headers)"}
+{"time":"2025-08-01T22:07:31.600992168+08:00","level":"INFO","msg":"api: retrying error","error":"Post \"https://api.bandw.top/files/gia0603yucca/stage2_08011616_2datasets_qweninstruct/790wkw0g/file_stream\": read tcp 10.1.11.100:49046->104.21.20.172:443: read: connection timed out"}
+{"time":"2025-08-01T22:09:05.97302798+08:00","level":"INFO","msg":"api: retrying error","error":"Post \"https://api.bandw.top/graphql\": context deadline exceeded"}
+{"time":"2025-08-01T22:09:46.459558104+08:00","level":"INFO","msg":"api: retrying error","error":"Post \"https://api.bandw.top/files/gia0603yucca/stage2_08011616_2datasets_qweninstruct/790wkw0g/file_stream\": read tcp 10.1.11.100:47316->104.21.20.172:443: read: connection reset by peer"}
+{"time":"2025-08-01T22:10:35.97491736+08:00","level":"INFO","msg":"api: retrying error","error":"Post \"https://api.bandw.top/graphql\": net/http: request canceled (Client.Timeout exceeded while awaiting headers)"}
+{"time":"2025-08-01T22:11:08.18377384+08:00","level":"INFO","msg":"api: retrying error","error":"Post \"https://api.bandw.top/graphql\": context deadline exceeded"}
+{"time":"2025-08-01T22:11:25.779600104+08:00","level":"INFO","msg":"api: retrying error","error":"Post \"https://api.bandw.top/files/gia0603yucca/stage2_08011616_2datasets_qweninstruct/790wkw0g/file_stream\": read tcp 10.1.11.100:34352->172.67.193.61:443: read: connection reset by peer"}
+{"time":"2025-08-01T22:11:42.262306077+08:00","level":"INFO","msg":"api: retrying error","error":"Post \"https://api.bandw.top/graphql\": context deadline exceeded"}
+{"time":"2025-08-01T22:12:50.976429474+08:00","level":"INFO","msg":"api: retrying error","error":"Post \"https://api.bandw.top/graphql\": net/http: request canceled (Client.Timeout exceeded while awaiting headers)"}
+{"time":"2025-08-01T22:13:23.229614619+08:00","level":"INFO","msg":"api: retrying error","error":"Post \"https://api.bandw.top/graphql\": context deadline exceeded"}
+{"time":"2025-08-01T22:13:57.892901227+08:00","level":"INFO","msg":"api: retrying error","error":"Post \"https://api.bandw.top/graphql\": net/http: request canceled (Client.Timeout exceeded while awaiting headers)"}
+{"time":"2025-08-01T22:15:04.820393306+08:00","level":"INFO","msg":"api: retrying error","error":"Post \"https://api.bandw.top/files/gia0603yucca/stage2_08011616_2datasets_qweninstruct/790wkw0g/file_stream\": read tcp 10.1.11.100:60344->104.21.20.172:443: read: connection reset by peer"}
+{"time":"2025-08-01T22:15:45.367773065+08:00","level":"INFO","msg":"api: retrying error","error":"Post \"https://api.bandw.top/files/gia0603yucca/stage2_08011616_2datasets_qweninstruct/790wkw0g/file_stream\": read tcp 10.1.11.100:54996->172.67.193.61:443: read: connection reset by peer"}
+{"time":"2025-08-01T22:17:57.605995597+08:00","level":"INFO","msg":"api: retrying error","error":"Post \"https://api.bandw.top/files/gia0603yucca/stage2_08011616_2datasets_qweninstruct/790wkw0g/file_stream\": unexpected EOF"}
+{"time":"2025-08-01T22:21:13.87299516+08:00","level":"INFO","msg":"api: retrying error","error":"Post \"https://api.bandw.top/files/gia0603yucca/stage2_08011616_2datasets_qweninstruct/790wkw0g/file_stream\": read tcp 10.1.11.100:49910->104.21.20.172:443: read: connection timed out"}
+{"time":"2025-08-01T22:24:16.144029987+08:00","level":"INFO","msg":"api: retrying error","error":"Post \"https://api.bandw.top/files/gia0603yucca/stage2_08011616_2datasets_qweninstruct/790wkw0g/file_stream\": read tcp 10.1.11.100:41940->172.67.193.61:443: read: connection timed out"}
+{"time":"2025-08-01T22:27:28.144021751+08:00","level":"INFO","msg":"api: retrying error","error":"Post \"https://api.bandw.top/files/gia0603yucca/stage2_08011616_2datasets_qweninstruct/790wkw0g/file_stream\": read tcp 10.1.11.100:46560->172.67.193.61:443: read: connection timed out"}
+{"time":"2025-08-01T22:28:35.984483121+08:00","level":"INFO","msg":"api: retrying error","error":"Post \"https://api.bandw.top/graphql\": net/http: request canceled (Client.Timeout exceeded while awaiting headers)"}
+{"time":"2025-08-01T22:31:45.168984421+08:00","level":"INFO","msg":"api: retrying error","error":"Post \"https://api.bandw.top/files/gia0603yucca/stage2_08011616_2datasets_qweninstruct/790wkw0g/file_stream\": read tcp 10.1.11.100:38656->172.67.193.61:443: read: connection timed out"}
+{"time":"2025-08-01T22:35:18.671993015+08:00","level":"INFO","msg":"api: retrying error","error":"Post \"https://api.bandw.top/files/gia0603yucca/stage2_08011616_2datasets_qweninstruct/790wkw0g/file_stream\": read tcp 10.1.11.100:49116->104.21.20.172:443: read: connection timed out"}
+{"time":"2025-08-01T22:39:02.928024828+08:00","level":"INFO","msg":"api: retrying error","error":"Post \"https://api.bandw.top/files/gia0603yucca/stage2_08011616_2datasets_qweninstruct/790wkw0g/file_stream\": read tcp 10.1.11.100:40350->172.67.193.61:443: read: connection timed out"}
+{"time":"2025-08-01T22:39:35.380477212+08:00","level":"INFO","msg":"api: retrying error","error":"Post \"https://api.bandw.top/files/gia0603yucca/stage2_08011616_2datasets_qweninstruct/790wkw0g/file_stream\": read tcp 10.1.11.100:35004->104.21.20.172:443: read: connection reset by peer"}
+{"time":"2025-08-01T22:39:55.432951061+08:00","level":"INFO","msg":"api: retrying error","error":"Post \"https://api.bandw.top/files/gia0603yucca/stage2_08011616_2datasets_qweninstruct/790wkw0g/file_stream\": unexpected EOF"}
+{"time":"2025-08-01T22:41:50.991964572+08:00","level":"INFO","msg":"api: retrying error","error":"Post \"https://api.bandw.top/graphql\": net/http: request canceled (Client.Timeout exceeded while awaiting headers)"}
+{"time":"2025-08-01T22:41:53.61821248+08:00","level":"INFO","msg":"api: retrying error","error":"Post \"https://api.bandw.top/files/gia0603yucca/stage2_08011616_2datasets_qweninstruct/790wkw0g/file_stream\": read tcp 10.1.11.100:37852->104.21.20.172:443: read: connection reset by peer"}
+{"time":"2025-08-01T22:43:35.932757702+08:00","level":"INFO","msg":"api: retrying error","error":"Post \"https://api.bandw.top/files/gia0603yucca/stage2_08011616_2datasets_qweninstruct/790wkw0g/file_stream\": http2: client conn is closed"}
+{"time":"2025-08-01T22:44:04.110612747+08:00","level":"INFO","msg":"api: retrying error","error":"Post \"https://api.bandw.top/files/gia0603yucca/stage2_08011616_2datasets_qweninstruct/790wkw0g/file_stream\": read tcp 10.1.11.100:53322->104.21.20.172:443: read: connection reset by peer"}
+{"time":"2025-08-01T22:46:35.994608155+08:00","level":"INFO","msg":"api: retrying error","error":"Post \"https://api.bandw.top/graphql\": context deadline exceeded"}
+{"time":"2025-08-01T22:47:08.18720141+08:00","level":"INFO","msg":"api: retrying error","error":"Post \"https://api.bandw.top/graphql\": net/http: request canceled (Client.Timeout exceeded while awaiting headers)"}
+{"time":"2025-08-01T22:47:34.92805984+08:00","level":"INFO","msg":"api: retrying error","error":"Post \"https://api.bandw.top/files/gia0603yucca/stage2_08011616_2datasets_qweninstruct/790wkw0g/file_stream\": read tcp 10.1.11.100:51190->172.67.193.61:443: read: connection timed out"}
+{"time":"2025-08-01T22:47:42.321052323+08:00","level":"INFO","msg":"api: retrying error","error":"Post \"https://api.bandw.top/graphql\": context deadline exceeded"}
+{"time":"2025-08-01T22:50:05.996775834+08:00","level":"INFO","msg":"api: retrying error","error":"Post \"https://api.bandw.top/graphql\": context deadline exceeded"}
+{"time":"2025-08-01T22:50:30.544005531+08:00","level":"INFO","msg":"api: retrying error","error":"Post \"https://api.bandw.top/files/gia0603yucca/stage2_08011616_2datasets_qweninstruct/790wkw0g/file_stream\": read tcp 10.1.11.100:49020->172.67.193.61:443: read: connection timed out"}
+{"time":"2025-08-01T22:50:38.303517763+08:00","level":"INFO","msg":"api: retrying error","error":"Post \"https://api.bandw.top/graphql\": net/http: request canceled (Client.Timeout exceeded while awaiting headers)"}
+{"time":"2025-08-01T22:51:12.580566425+08:00","level":"INFO","msg":"api: retrying error","error":"Post \"https://api.bandw.top/graphql\": context deadline exceeded"}
+{"time":"2025-08-01T22:51:51.301848016+08:00","level":"INFO","msg":"api: retrying error","error":"Post \"https://api.bandw.top/graphql\": context deadline exceeded"}
+{"time":"2025-08-01T22:52:39.862980102+08:00","level":"INFO","msg":"api: retrying error","error":"Post \"https://api.bandw.top/graphql\": context deadline exceeded"}
+{"time":"2025-08-01T22:53:43.460661511+08:00","level":"INFO","msg":"api: retrying error","error":"Post \"https://api.bandw.top/graphql\": net/http: request canceled (Client.Timeout exceeded while awaiting headers)"}
+{"time":"2025-08-01T22:55:13.462771183+08:00","level":"INFO","msg":"api: retrying error","error":"Post \"https://api.bandw.top/graphql\": context deadline exceeded"}
+{"time":"2025-08-01T22:55:50.033040855+08:00","level":"INFO","msg":"api: retrying error","error":"Post \"https://api.bandw.top/files/gia0603yucca/stage2_08011616_2datasets_qweninstruct/790wkw0g/file_stream\": read tcp 10.1.11.100:38370->172.67.193.61:443: read: connection timed out"}
+{"time":"2025-08-01T22:56:24.977146749+08:00","level":"INFO","msg":"api: retrying error","error":"Post \"https://api.bandw.top/files/gia0603yucca/stage2_08011616_2datasets_qweninstruct/790wkw0g/file_stream\": read tcp 10.1.11.100:39018->104.21.20.172:443: read: connection reset by peer"}
+{"time":"2025-08-01T22:56:43.465087253+08:00","level":"INFO","msg":"api: retrying error","error":"Post \"https://api.bandw.top/graphql\": net/http: request canceled (Client.Timeout exceeded while awaiting headers)"}
+{"time":"2025-08-01T22:58:13.466056502+08:00","level":"INFO","msg":"api: retrying error","error":"Post \"https://api.bandw.top/graphql\": context deadline exceeded"}
+{"time":"2025-08-01T22:59:29.003634152+08:00","level":"INFO","msg":"api: retrying error","error":"Post \"https://api.bandw.top/files/gia0603yucca/stage2_08011616_2datasets_qweninstruct/790wkw0g/file_stream\": context deadline exceeded (Client.Timeout exceeded while awaiting headers)"}
+{"time":"2025-08-01T22:59:35.997189916+08:00","level":"WARN","msg":"sender: taking a long time","seconds":600.000438436,"work":"WorkRecord(*service_go_proto.Request_StopStatus); Control(local:true mailbox_slot:\"zbp0fazrv773\" connection_id:\"127.0.0.1:32880\")"}
+{"time":"2025-08-01T22:59:43.468121454+08:00","level":"INFO","msg":"api: retrying error","error":"Post \"https://api.bandw.top/graphql\": net/http: request canceled (Client.Timeout exceeded while awaiting headers)"}
+{"time":"2025-08-01T22:59:56.405426213+08:00","level":"INFO","msg":"api: retrying error","error":"Post \"https://api.bandw.top/files/gia0603yucca/stage2_08011616_2datasets_qweninstruct/790wkw0g/file_stream\": unexpected EOF"}
+{"time":"2025-08-01T22:59:56.885071194+08:00","level":"WARN","msg":"runwork: taking a long time","seconds":600.000154693,"work":"WorkRecord(*service_go_proto.Record_OutputRaw); Control(connection_id:\"127.0.0.1:32880\")"}
+{"time":"2025-08-01T23:00:05.004711621+08:00","level":"WARN","msg":"runwork: taking a long time","seconds":600.00033328,"work":"WorkRecord(*service_go_proto.Record_Stats); Control(always_send:true)"}
+{"time":"2025-08-01T23:00:05.076924009+08:00","level":"WARN","msg":"runwork: taking a long time","seconds":600.000302659,"work":"WorkRecord(*service_go_proto.Record_Stats); Control(always_send:true)"}
+{"time":"2025-08-01T23:01:13.46903123+08:00","level":"INFO","msg":"api: retrying error","error":"Post \"https://api.bandw.top/graphql\": context deadline exceeded"}
+{"time":"2025-08-01T23:02:43.472282526+08:00","level":"INFO","msg":"sender: succeeded after taking longer than expected","seconds":787.475510948,"work":"WorkRecord(*service_go_proto.Request_StopStatus); Control(local:true mailbox_slot:\"zbp0fazrv773\" connection_id:\"127.0.0.1:32880\")"}
+{"time":"2025-08-01T23:02:43.472384259+08:00","level":"INFO","msg":"runwork: succeeded after taking longer than expected","seconds":758.39578933,"work":"WorkRecord(*service_go_proto.Record_Stats); Control(always_send:true)"}
+{"time":"2025-08-01T23:02:43.472386124+08:00","level":"INFO","msg":"runwork: succeeded after taking longer than expected","seconds":758.468026974,"work":"WorkRecord(*service_go_proto.Record_Stats); Control(always_send:true)"}
+{"time":"2025-08-01T23:02:43.47239327+08:00","level":"INFO","msg":"runwork: succeeded after taking longer than expected","seconds":766.587479442,"work":"WorkRecord(*service_go_proto.Record_OutputRaw); Control(connection_id:\"127.0.0.1:32880\")"}
+{"time":"2025-08-01T23:03:03.192446429+08:00","level":"INFO","msg":"api: retrying error","error":"Post \"https://api.bandw.top/files/gia0603yucca/stage2_08011616_2datasets_qweninstruct/790wkw0g/file_stream\": read tcp 10.1.11.100:44012->172.67.193.61:443: read: connection reset by peer"}
+{"time":"2025-08-01T23:03:21.005846988+08:00","level":"INFO","msg":"api: retrying error","error":"Post \"https://api.bandw.top/graphql\": context deadline exceeded"}
+{"time":"2025-08-01T23:08:40.592045526+08:00","level":"INFO","msg":"api: retrying error","error":"Post \"https://api.bandw.top/files/gia0603yucca/stage2_08011616_2datasets_qweninstruct/790wkw0g/file_stream\": read tcp 10.1.11.100:38310->104.21.20.172:443: read: connection timed out"}
+{"time":"2025-08-01T23:11:46.960031857+08:00","level":"INFO","msg":"api: retrying error","error":"Post \"https://api.bandw.top/files/gia0603yucca/stage2_08011616_2datasets_qweninstruct/790wkw0g/file_stream\": read tcp 10.1.11.100:55084->104.21.20.172:443: read: connection timed out"}
+{"time":"2025-08-01T23:12:24.849142697+08:00","level":"INFO","msg":"api: retrying error","error":"Post \"https://api.bandw.top/files/gia0603yucca/stage2_08011616_2datasets_qweninstruct/790wkw0g/file_stream\": read tcp 10.1.11.100:49670->104.21.20.172:443: read: connection reset by peer"}
+{"time":"2025-08-01T23:15:08.177028412+08:00","level":"INFO","msg":"api: retrying error","error":"Post \"https://api.bandw.top/files/gia0603yucca/stage2_08011616_2datasets_qweninstruct/790wkw0g/file_stream\": read tcp 10.1.11.100:52862->104.21.20.172:443: read: connection timed out"}
+{"time":"2025-08-01T23:16:46.602539365+08:00","level":"INFO","msg":"api: retrying error","error":"Post \"https://api.bandw.top/files/gia0603yucca/stage2_08011616_2datasets_qweninstruct/790wkw0g/file_stream\": read tcp 10.1.11.100:45510->104.21.20.172:443: read: connection reset by peer"}
+{"time":"2025-08-01T23:18:23.28263426+08:00","level":"INFO","msg":"api: retrying error","error":"Post \"https://api.bandw.top/graphql\": read tcp 10.1.11.100:48380->172.67.193.61:443: read: connection reset by peer"}
+{"time":"2025-08-01T23:19:27.841519011+08:00","level":"INFO","msg":"api: retrying error","error":"Post \"https://api.bandw.top/files/gia0603yucca/stage2_08011616_2datasets_qweninstruct/790wkw0g/file_stream\": unexpected EOF"}
+{"time":"2025-08-01T23:23:01.165294384+08:00","level":"INFO","msg":"api: retrying error","error":"Post \"https://api.bandw.top/files/gia0603yucca/stage2_08011616_2datasets_qweninstruct/790wkw0g/file_stream\": context deadline exceeded"}
+{"time":"2025-08-01T23:26:53.200081773+08:00","level":"INFO","msg":"api: retrying error","error":"Post \"https://api.bandw.top/files/gia0603yucca/stage2_08011616_2datasets_qweninstruct/790wkw0g/file_stream\": read tcp 10.1.11.100:43360->104.21.20.172:443: read: connection timed out"}
+{"time":"2025-08-01T23:34:57.039993804+08:00","level":"INFO","msg":"api: retrying error","error":"Post \"https://api.bandw.top/files/gia0603yucca/stage2_08011616_2datasets_qweninstruct/790wkw0g/file_stream\": read tcp 10.1.11.100:52350->172.67.193.61:443: read: connection timed out"}
+{"time":"2025-08-01T23:38:16.720010379+08:00","level":"INFO","msg":"api: retrying error","error":"Post \"https://api.bandw.top/files/gia0603yucca/stage2_08011616_2datasets_qweninstruct/790wkw0g/file_stream\": read tcp 10.1.11.100:57172->104.21.20.172:443: read: connection timed out"}
+{"time":"2025-08-01T23:41:21.040008763+08:00","level":"INFO","msg":"api: retrying error","error":"Post \"https://api.bandw.top/files/gia0603yucca/stage2_08011616_2datasets_qweninstruct/790wkw0g/file_stream\": read tcp 10.1.11.100:36508->104.21.20.172:443: read: connection timed out"}
+{"time":"2025-08-01T23:44:44.816004249+08:00","level":"INFO","msg":"api: retrying error","error":"Post \"https://api.bandw.top/files/gia0603yucca/stage2_08011616_2datasets_qweninstruct/790wkw0g/file_stream\": read tcp 10.1.11.100:55942->172.67.193.61:443: read: connection timed out"}
+{"time":"2025-08-01T23:48:52.62500069+08:00","level":"INFO","msg":"api: retrying error","error":"Post \"https://api.bandw.top/files/gia0603yucca/stage2_08011616_2datasets_qweninstruct/790wkw0g/file_stream\": read tcp 10.1.11.100:47172->172.67.193.61:443: read: connection timed out"}
+{"time":"2025-08-01T23:51:31.345006101+08:00","level":"INFO","msg":"api: retrying error","error":"Post \"https://api.bandw.top/files/gia0603yucca/stage2_08011616_2datasets_qweninstruct/790wkw0g/file_stream\": read tcp 10.1.11.100:52938->104.21.20.172:443: read: connection timed out"}
+{"time":"2025-08-01T23:54:39.760994652+08:00","level":"INFO","msg":"api: retrying error","error":"Post \"https://api.bandw.top/files/gia0603yucca/stage2_08011616_2datasets_qweninstruct/790wkw0g/file_stream\": read tcp 10.1.11.100:57488->104.21.20.172:443: read: connection timed out"}
+{"time":"2025-08-01T23:56:53.243171871+08:00","level":"INFO","msg":"api: retrying error","error":"Post \"https://api.bandw.top/files/gia0603yucca/stage2_08011616_2datasets_qweninstruct/790wkw0g/file_stream\": unexpected EOF"}
+{"time":"2025-08-02T00:01:11.439949517+08:00","level":"INFO","msg":"api: retrying error","error":"Post \"https://api.bandw.top/files/gia0603yucca/stage2_08011616_2datasets_qweninstruct/790wkw0g/file_stream\": read tcp 10.1.11.100:48808->172.67.193.61:443: read: connection timed out"}
+{"time":"2025-08-02T00:03:57.840030605+08:00","level":"INFO","msg":"api: retrying error","error":"Post \"https://api.bandw.top/files/gia0603yucca/stage2_08011616_2datasets_qweninstruct/790wkw0g/file_stream\": read tcp 10.1.11.100:58796->172.67.193.61:443: read: connection timed out"}
+{"time":"2025-08-02T00:04:51.340371717+08:00","level":"INFO","msg":"api: retrying error","error":"Post \"https://api.bandw.top/graphql\": read tcp 10.1.11.100:40756->104.21.20.172:443: read: connection reset by peer"}
+{"time":"2025-08-02T00:06:51.040127189+08:00","level":"INFO","msg":"api: retrying error","error":"Post \"https://api.bandw.top/graphql\": net/http: request canceled (Client.Timeout exceeded while awaiting headers)"}
+{"time":"2025-08-02T00:12:23.188896321+08:00","level":"INFO","msg":"api: retrying error","error":"Post \"https://api.bandw.top/graphql\": read tcp 10.1.11.100:52718->172.67.193.61:443: read: connection reset by peer"}
+{"time":"2025-08-02T00:19:09.499959872+08:00","level":"INFO","msg":"api: retrying error","error":"Post \"https://api.bandw.top/files/gia0603yucca/stage2_08011616_2datasets_qweninstruct/790wkw0g/file_stream\": unexpected EOF"}
+{"time":"2025-08-02T00:27:17.136029213+08:00","level":"INFO","msg":"api: retrying error","error":"Post \"https://api.bandw.top/files/gia0603yucca/stage2_08011616_2datasets_qweninstruct/790wkw0g/file_stream\": read tcp 10.1.11.100:33250->104.21.20.172:443: read: connection timed out"}
+{"time":"2025-08-02T00:31:24.717936901+08:00","level":"INFO","msg":"api: retrying error","error":"Post \"https://api.bandw.top/files/gia0603yucca/stage2_08011616_2datasets_qweninstruct/790wkw0g/file_stream\": read tcp 10.1.11.100:42528->104.21.20.172:443: read: connection reset by peer"}
+{"time":"2025-08-02T00:49:10.92803379+08:00","level":"INFO","msg":"api: retrying error","error":"Post \"https://api.bandw.top/files/gia0603yucca/stage2_08011616_2datasets_qweninstruct/790wkw0g/file_stream\": read tcp 10.1.11.100:49640->172.67.193.61:443: read: connection timed out"}
+{"time":"2025-08-02T00:53:16.688059564+08:00","level":"INFO","msg":"api: retrying error","error":"Post \"https://api.bandw.top/files/gia0603yucca/stage2_08011616_2datasets_qweninstruct/790wkw0g/file_stream\": read tcp 10.1.11.100:41832->172.67.193.61:443: read: connection timed out"}
+{"time":"2025-08-02T01:03:03.440032369+08:00","level":"INFO","msg":"api: retrying error","error":"Post \"https://api.bandw.top/files/gia0603yucca/stage2_08011616_2datasets_qweninstruct/790wkw0g/file_stream\": read tcp 10.1.11.100:36398->104.21.20.172:443: read: connection timed out"}
+{"time":"2025-08-02T01:15:25.202630664+08:00","level":"INFO","msg":"api: retrying error","error":"Post \"https://api.bandw.top/files/gia0603yucca/stage2_08011616_2datasets_qweninstruct/790wkw0g/file_stream\": read tcp 10.1.11.100:35078->104.21.20.172:443: read: connection reset by peer"}
+{"time":"2025-08-02T01:23:06.089931169+08:00","level":"INFO","msg":"api: retrying error","error":"Post \"https://api.bandw.top/graphql\": net/http: request canceled (Client.Timeout exceeded while awaiting headers)"}
+{"time":"2025-08-02T01:28:16.400050797+08:00","level":"INFO","msg":"api: retrying error","error":"Post \"https://api.bandw.top/files/gia0603yucca/stage2_08011616_2datasets_qweninstruct/790wkw0g/file_stream\": read tcp 10.1.11.100:37336->172.67.193.61:443: read: connection timed out"}
+{"time":"2025-08-02T02:22:06.219804028+08:00","level":"INFO","msg":"api: retrying error","error":"Post \"https://api.bandw.top/graphql\": net/http: request canceled (Client.Timeout exceeded while awaiting headers)"}
+{"time":"2025-08-02T03:05:52.144050294+08:00","level":"INFO","msg":"api: retrying error","error":"Post \"https://api.bandw.top/files/gia0603yucca/stage2_08011616_2datasets_qweninstruct/790wkw0g/file_stream\": read tcp 10.1.11.100:51464->172.67.193.61:443: read: connection timed out"}
+{"time":"2025-08-02T06:13:41.315904777+08:00","level":"INFO","msg":"api: retrying HTTP error","status":504,"url":"https://api.bandw.top/files/gia0603yucca/stage2_08011616_2datasets_qweninstruct/790wkw0g/file_stream","body":"error code: 504"}
+{"time":"2025-08-02T06:13:52.289991198+08:00","level":"INFO","msg":"api: retrying error","error":"Post \"https://api.bandw.top/graphql\": context deadline exceeded"}
+{"time":"2025-08-02T06:55:22.310503234+08:00","level":"INFO","msg":"api: retrying error","error":"Post \"https://api.bandw.top/graphql\": context deadline exceeded"}
+{"time":"2025-08-02T06:55:25.73261549+08:00","level":"INFO","msg":"api: retrying HTTP error","status":504,"url":"https://api.bandw.top/files/gia0603yucca/stage2_08011616_2datasets_qweninstruct/790wkw0g/file_stream","body":"error code: 504"}
+{"time":"2025-08-02T07:13:22.52185475+08:00","level":"INFO","msg":"api: retrying error","error":"Post \"https://api.bandw.top/graphql\": read tcp 10.1.11.100:41606->104.21.20.172:443: read: connection reset by peer"}
+{"time":"2025-08-02T07:38:52.528290374+08:00","level":"INFO","msg":"api: retrying error","error":"Post \"https://api.bandw.top/graphql\": read tcp 10.1.11.100:38352->104.21.20.172:443: read: connection reset by peer"}
+{"time":"2025-08-02T09:31:16.192630538+08:00","level":"INFO","msg":"stream: closing","id":"790wkw0g"}
+{"time":"2025-08-02T09:31:16.192679202+08:00","level":"INFO","msg":"Stopping system monitor"}
+{"time":"2025-08-02T09:31:16.194392232+08:00","level":"INFO","msg":"Stopped system monitor"}
+{"time":"2025-08-02T09:31:21.149902109+08:00","level":"INFO","msg":"fileTransfer: Close: file transfer manager closed"}
+{"time":"2025-08-02T09:31:22.465229216+08:00","level":"INFO","msg":"handler: closed","stream_id":"790wkw0g"}
+{"time":"2025-08-02T09:31:22.465258111+08:00","level":"INFO","msg":"writer: Close: closed","stream_id":"790wkw0g"}
+{"time":"2025-08-02T09:31:22.465263174+08:00","level":"INFO","msg":"sender: closed","stream_id":"790wkw0g"}
+{"time":"2025-08-02T09:31:22.469479316+08:00","level":"INFO","msg":"stream: closed","id":"790wkw0g"}

all_checkpoints/stage2_08011616_2datasets_withoutpretrain/wandb/run-20250801_162632-790wkw0g/logs/debug.log

@@ -0,0 +1,24 @@
+2025-08-01 16:26:32,237 INFO    MainThread:764673 [wandb_setup.py:_flush():70] Current SDK version is 0.19.11
+2025-08-01 16:26:32,237 INFO    MainThread:764673 [wandb_setup.py:_flush():70] Configure stats pid to 764673
+2025-08-01 16:26:32,237 INFO    MainThread:764673 [wandb_setup.py:_flush():70] Loading settings from /root/.config/wandb/settings
+2025-08-01 16:26:32,237 INFO    MainThread:764673 [wandb_setup.py:_flush():70] Loading settings from /nas/shared/kilab/wangyujia/ProtT3/wandb/settings
+2025-08-01 16:26:32,237 INFO    MainThread:764673 [wandb_setup.py:_flush():70] Loading settings from environment variables
+2025-08-01 16:26:32,237 INFO    MainThread:764673 [wandb_init.py:setup_run_log_directory():724] Logging user logs to ./all_checkpoints/stage2_08011616_2datasets_qweninstruct/wandb/run-20250801_162632-790wkw0g/logs/debug.log
+2025-08-01 16:26:32,237 INFO    MainThread:764673 [wandb_init.py:setup_run_log_directory():725] Logging internal logs to ./all_checkpoints/stage2_08011616_2datasets_qweninstruct/wandb/run-20250801_162632-790wkw0g/logs/debug-internal.log
+2025-08-01 16:26:32,237 INFO    MainThread:764673 [wandb_init.py:init():852] calling init triggers
+2025-08-01 16:26:32,237 INFO    MainThread:764673 [wandb_init.py:init():857] wandb.init called with sweep_config: {}
+config: {'_wandb': {}}
+2025-08-01 16:26:32,237 INFO    MainThread:764673 [wandb_init.py:init():893] starting backend
+2025-08-01 16:26:32,237 INFO    MainThread:764673 [wandb_init.py:init():897] sending inform_init request
+2025-08-01 16:26:32,241 INFO    MainThread:764673 [backend.py:_multiprocessing_setup():101] multiprocessing start_methods=fork,spawn,forkserver, using: spawn
+2025-08-01 16:26:32,241 INFO    MainThread:764673 [wandb_init.py:init():907] backend started and connected
+2025-08-01 16:26:32,242 INFO    MainThread:764673 [wandb_init.py:init():1005] updated telemetry
+2025-08-01 16:26:32,306 INFO    MainThread:764673 [wandb_init.py:init():1029] communicating run to backend with 90.0 second timeout
+2025-08-01 16:26:49,938 INFO    MainThread:764673 [wandb_init.py:init():1104] starting run threads in backend
+2025-08-01 16:26:50,132 INFO    MainThread:764673 [wandb_run.py:_console_start():2573] atexit reg
+2025-08-01 16:26:50,132 INFO    MainThread:764673 [wandb_run.py:_redirect():2421] redirect: wrap_raw
+2025-08-01 16:26:50,132 INFO    MainThread:764673 [wandb_run.py:_redirect():2490] Wrapping output streams.
+2025-08-01 16:26:50,132 INFO    MainThread:764673 [wandb_run.py:_redirect():2513] Redirects installed.
+2025-08-01 16:26:50,134 INFO    MainThread:764673 [wandb_init.py:init():1150] run started, returning control to user process
+2025-08-01 16:26:58,724 INFO    MainThread:764673 [wandb_run.py:_config_callback():1436] config_cb None None {'filename': 'stage2_08011616_2datasets_qweninstruct', 'seed': 42, 'mode': 'train', 'strategy': 'deepspeed', 'accelerator': 'gpu', 'devices': '0,1,2,3,4,5,6,7', 'precision': 'bf16-mixed', 'max_epochs': 10, 'accumulate_grad_batches': 1, 'check_val_every_n_epoch': 1, 'enable_flash': False, 'use_wandb_logger': True, 'mix_dataset': False, 'dataset': 'swiss-prot', 'save_every_n_epochs': 2, 'bert_name': '/nas/shared/kilab/wangyujia/ProtT3/plm_model/microsoft', 'cross_attention_freq': 2, 'num_query_token': 8, 'qformer_tune': 'train', 'llm_name': '/nas/shared/kilab/hf-hub/models--Qwen--Qwen2.5-7B-Instruct/snapshots/a09a35458c702b33eeacc393d103063234e8bc28', 'num_beams': 5, 'do_sample': False, 'max_inference_len': 128, 'min_inference_len': 1, 'llm_tune': 'mid_lora', 'peft_config': '', 'peft_dir': '', 'plm_model': '/nas/shared/kilab/wangyujia/ProtT3/plm_model/esm2-150m', 'plm_tune': 'freeze', 'lora_r': 8, 'lora_alpha': 16, 'lora_dropout': 0.1, 'enbale_gradient_checkpointing': False, 'weight_decay': 0.05, 'init_lr': 0.0001, 'min_lr': 1e-05, 'warmup_lr': 1e-06, 'warmup_steps': 1000, 'lr_decay_rate': 0.9, 'scheduler': 'linear_warmup_cosine_lr', 'stage1_path': '/nas/shared/kilab/wangyujia/ProtT3/all_checkpoints/stage1_07041727_2dataset/epoch=29.ckpt/converted.ckpt', 'stage2_path': '', 'init_checkpoint': '/nas/shared/kilab/wangyujia/ProtT3/all_checkpoints/stage2_07070513_2datasets_construct/epoch=09.ckpt/converted.ckpt', 'caption_eval_epoch': 5, 'num_workers': 8, 'batch_size': 4, 'inference_batch_size': 4, 'root': 'data', 'text_max_len': 2048, 'q_max_len': 29, 'a_max_len': 36, 'prot_max_len': 1024, 'prompt': 'The protein has the following properties:', 'filter_side_qa': False}
+2025-08-02 09:31:16,189 INFO    MsgRouterThr:764673 [mailbox.py:close():129] [no run ID] Closing mailbox, abandoning 1 handles.

all_checkpoints/stage2_08011616_2datasets_withoutpretrain/wandb/run-20250801_162632-790wkw0g/run-790wkw0g.wandb

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:cb5d0be9388288b74f1f5a6f3f8c39d9ee6aa7643fe0c3426a6f51210ad77dc1
+size 84438700

data/.gitignore

@@ -0,0 +1,2 @@
+*
+!.gitignore

data/OntoProteinDatasetV2/test.txt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:87c0edf3fd59defd24decb12c8af64d9bb5fa1ef727fce2bf5fc36c06fe4066f
+size 12085289

data/OntoProteinDatasetV2/train.txt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:632a5830aea6b3029feaf2eeba5a1c2b63e2badc61b72f92734c517aef85e879
+size 478799520

data/OntoProteinDatasetV2/valid.txt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:40cefeb9b08cff978fee2ce21f259974fa8881bbe6614f04e896ca7560d4e11f
+size 11829916

data/PDBDataset/abstract.json

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:59cf438e3325c38809ab41fc09fe00ec78fc4bd41d7073f0c22f664168ea5ca3
+size 190905202

data/PDBDataset/q_types.txt

@@ -0,0 +1,30 @@
+What is the nucleic acid polymer entity type for this protein?	String structure/property
+When is this protein first published?	Number side information
+How many polymer monomers does this protein have?	Number structure/property
+How many assemblies does this protein have?	Number structure/property
+How many heavy solvent atom coordinates records does this protein have?	Number side information
+Does this protein have cis-peptide linkages?	String structure/property
+Does this protein contain branched entities?	String structure/property
+How many entities does this protein have?	Number structure/property
+Does this protein contain solvent entities?	String structure/property
+What is the polymer entity type for this protein?	String structure/property
+How many nucleic acid polymer entities (DNA or RNA) does this protein have?	Number structure/property
+What is the polymer entity composition for this protein?	String structure/property
+Does this protein contain DNA polymer entities?	String structure/property
+Does this protein contain non-polymer entities?	String structure/property
+How many heavy atom coordinates records does this protein have?	Number side information
+How many intermolecular metalic bonds does this protein have?	Number structure/property
+Does this protein have hybrid nucleic acid polymer entities?	String structure/property
+What is the molecular mass (KDa) of polymer and non-polymer entities (exclusive of solvent) for this protein?	Number structure/property
+Is this protein determined by experimental or computational methods?	String side information
+What is the radiation wavelength in angstroms for this protein?	Number structure/property
+Does this protein have unmodeled polymer monomers?	String side information
+How many intermolecular covalent bonds does this protein have?	Number structure/property
+How many model structures deposited for this protein?	Number side information
+What are the software programs reported in connection with the production of this protein?	String side information
+How many hydrogen atom coordinates records does this protein have?	Number side information
+What experimental method(s) were used to determine the structure of this protein?	String side information
+Does this protein contain polymer entities?	String structure/property
+Does this protein contain RNA polymer entities?	String structure/property
+What are the bound nonpolymer components for this protein	String structure/property
+What are the terms characterizing the protein?	String structure/property

data/PDBDataset/qa_all.json

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:dfd9d775036717127b66aa9ba9c9d61f8fc324c4ea8fb548d6a1fd2ab4506194
+size 523476104

data/PDBDataset/train.txt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:494c6c0680bdedd58dc0da2df11393b123ad5a4998c68d432da0b1db777c8870
+size 34040909

data/SwissProtV3/train_set.jsonl

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:d92b283a50056b185af292082ea3c733d2c979e70ca233c19d4b06bf81381713
+size 312044762

data/protein-molecule/protein-text.zip

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:5ae14ae1857962332f2b73a3f7ef5651541e7be7d886bb71031b4fc745626f3a
+size 1566756112

data/protein-text/eval_assist.zipg3ebgjl7.tmp

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:28e38b555c6d456da8e2b95dd2af19f884b2b86687f733d745b7972079dff708
+size 113836360

data/protein-text/eval_assist.ziphwjr8q2y.tmp

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:58bd802e8fc350fd425e9a827245662790580004bc064aa84b265cad8f2e6d3f
+size 9836668

data/protein-text/eval_assist.zipzh1pdmj_.tmp

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:3ddebab3f8ea6218c001053e6861781b90b2b70ef429bf9239761c2c0c7927e0
+size 34493776

data_provider/bindingdb.py

@@ -0,0 +1,62 @@
+import json
+from pytorch_lightning import LightningDataModule
+from torch.utils.data import DataLoader, ConcatDataset, Dataset
+from data_provider.stage1_dm import SwissProtDataset, OntoProteinDataset
+import pandas as pd
+
+class BindingDB(Dataset):
+    def __init__(self, data_path, prompt='', return_prompt=False):
+        super(BindingDB, self).__init__()
+        self.data_path = data_path
+        self.user_prompt = prompt
+        self.return_prompt = return_prompt
+        
+        self.data_list = self._load_and_preprocess(self.data_path)
+        self.text2id = self._build_text_vocab()
+
+    def _load_and_preprocess(self, data_path):
+        data_list = []
+        df = pd.read_csv(data_path)
+        for _, row in df.iterrows():
+            try:
+                ligand_smiles = str(row['ligand']).strip()
+                prot_seq = str(row['protein']).strip()
+                result = str(row['ic50']).strip()
+
+                text_seq = f"<answer>{result}</answer>\n"
+ 
+                prompt = f"""
+                【Protein sequence (1-letter amino acid codes)】；{ligand_smiles}【Ligand structure (SMILES)】
+                Task: Evaluate the inhibitory effect of the ligand on the given protein.
+                Note: IC50 (half maximal inhibitory concentration) is the concentration of a substance required to inhibit 50% of the protein's activity. Lower IC50 values indicate stronger inhibition.
+                Based on the provided protein and ligand, predict the inhibitory strength by classifying the IC50 level:  
+                """
+                if self.user_prompt:
+                    prompt += self.user_prompt
+
+                # extra可以返回原始feather字符串，也可以返回feather_vals
+                 # 或 feather_raw
+                data_list.append((prot_seq, text_seq, prompt))
+            except Exception as e:
+                print(f"警告: 跳过有问题的行: {row}，原因: {e}")
+        return data_list
+
+    def _build_text_vocab(self):
+        text2id = {}
+        for _, text_seq, _ in self.data_list:
+            if text_seq not in text2id:
+                text2id[text_seq] = len(text2id)
+        return text2id
+
+    def shuffle(self):
+        random.shuffle(self.data_list)
+        return self
+
+    def __len__(self):
+        return len(self.data_list)
+
+    def __getitem__(self, index):
+        prot_seq, text_seq, prompt = self.data_list[index]
+        if self.return_prompt:
+            return prot_seq, prompt, text_seq,index
+        return prot_seq, text_seq, index

data_provider/gal_helpers.py

@@ -0,0 +1,45 @@
+import re
+
+
+# we split individual characters inside special tokens like [START_DNA]
+CUSTOM_SEQ_RE = re.compile(r"(\[START_(DNA|SMILES|I_SMILES|AMINO)])(.*?)(\[END_\2])")
+
+# token added to implement a custom sequence tokenization. This token is added at
+# corpus cleaning step and removed in pretokenization. The digits are added to increase the chance
+# that they do not occur in the corpus. The digits are escaped so that the token does not appear
+# literally in the source code in case we ever include it in the training data.
+SPLIT_MARKER = f"SPL{1}T-TH{1}S-Pl3A5E"
+
+def _insert_split_marker(m: re.Match):
+    """
+    Applies split marker based on a regex match of special tokens such as
+    [START_DNA].
+
+    Parameters
+    ----------
+    n : str
+        Input text to split
+
+    Returns
+    ----------
+    str - the text with the split token added
+    """
+    start_token, _, sequence, end_token = m.groups()
+    sequence = re.sub(r"(.)", fr"{SPLIT_MARKER}\1", sequence, flags=re.DOTALL)
+    return f"{start_token}{sequence}{SPLIT_MARKER}{end_token}"
+
+
+def escape_custom_split_sequence(text):
+    """
+    Applies custom splitting to the text for GALILEO's tokenization
+
+    Parameters
+    ----------
+    text : str
+        Input text to split
+
+    Returns
+    ----------
+    str - the text with the split token added
+    """
+    return CUSTOM_SEQ_RE.sub(_insert_split_marker, text)

data_provider/go.py

@@ -0,0 +1,237 @@
+import json
+from pytorch_lightning import LightningDataModule
+from torch.utils.data import DataLoader, ConcatDataset, Dataset
+from data_provider.stage1_dm import SwissProtDataset, OntoProteinDataset
+import pandas as pd
+
+class GO_BP(Dataset):
+    def __init__(self, data_path, prompt='', return_prompt=False):
+        super(GO_BP, self).__init__()
+        self.data_path = data_path
+        self.user_prompt = prompt
+        self.return_prompt = return_prompt
+        
+        self.data_list = self._load_and_preprocess(self.data_path)
+        self.text2id = self._build_text_vocab()
+
+    def _load_and_preprocess(self, data_path):
+        data_list = []
+        df = pd.read_csv(data_path)
+        for _, row in df.iterrows():
+            try:
+                prot_seq = str(row['question']).strip()
+                result = str(row['answer']).strip()
+
+                text_seq = f"<answer>{result}</answer>\n"
+ 
+                prompt = f"""
+【Task】Predict the biological processes involving the given protein.
+【Background】Each process is represented by a GO-BP term (e.g., GO:0008150), describing a series of molecular events relevant to protein function.
+【Output Format】List the predicted GO-BP terms, separated by commas, and wrap them in <answer> </answer> tags.  
+Example: <answer>GO:0008150, GO:0009987, GO:0050896</answer>
+"""
+                if self.user_prompt:
+                    prompt += self.user_prompt
+
+                # extra可以返回原始feather字符串，也可以返回feather_vals
+                 # 或 feather_raw
+                data_list.append((prot_seq, text_seq, prompt))
+            except Exception as e:
+                print(f"警告: 跳过有问题的行: {row}，原因: {e}")
+        return data_list
+
+    def _build_text_vocab(self):
+        text2id = {}
+        for _, text_seq, _ in self.data_list:
+            if text_seq not in text2id:
+                text2id[text_seq] = len(text2id)
+        return text2id
+
+    def shuffle(self):
+        random.shuffle(self.data_list)
+        return self
+
+    def __len__(self):
+        return len(self.data_list)
+
+    def __getitem__(self, index):
+        prot_seq, text_seq, prompt = self.data_list[index]
+        if self.return_prompt:
+            return prot_seq, prompt, text_seq,index
+        return prot_seq, text_seq, index
+
+
+class GO_CC(Dataset):
+    def __init__(self, data_path, prompt='', return_prompt=False):
+        super(GO_CC, self).__init__()
+        self.data_path = data_path
+        self.user_prompt = prompt
+        self.return_prompt = return_prompt
+        
+        self.data_list = self._load_and_preprocess(self.data_path)
+        self.text2id = self._build_text_vocab()
+
+    def _load_and_preprocess(self, data_path):
+        data_list = []
+        df = pd.read_csv(data_path)
+        for _, row in df.iterrows():
+            try:
+                prot_seq = str(row['question']).strip()
+                result = str(row['answer']).strip()
+
+                text_seq = f"<answer>{result}</answer>\n"
+ 
+                prompt = f"""
+【Task】Predict the cellular components associated with this protein.
+【Background】Each location is represented by a GO-CC term (e.g., GO:0005737), indicating where the protein functions within the cell.
+【Output Format】List the predicted GO-CC terms, separated by commas, and wrap them in <answer> </answer> tags.  
+Example: <answer>GO:0005737, GO:0005829, GO:0005886</answer>
+"""
+                if self.user_prompt:
+                    prompt += self.user_prompt
+
+                # extra可以返回原始feather字符串，也可以返回feather_vals
+                 # 或 feather_raw
+                data_list.append((prot_seq, text_seq, prompt))
+            except Exception as e:
+                print(f"警告: 跳过有问题的行: {row}，原因: {e}")
+        return data_list
+
+    def _build_text_vocab(self):
+        text2id = {}
+        for _, text_seq, _ in self.data_list:
+            if text_seq not in text2id:
+                text2id[text_seq] = len(text2id)
+        return text2id
+
+    def shuffle(self):
+        random.shuffle(self.data_list)
+        return self
+
+    def __len__(self):
+        return len(self.data_list)
+
+    def __getitem__(self, index):
+        prot_seq, text_seq, prompt = self.data_list[index]
+        if self.return_prompt:
+            return prot_seq, prompt, text_seq,index
+        return prot_seq, text_seq, index
+
+
+class GO_MF(Dataset):
+    def __init__(self, data_path, prompt='', return_prompt=False):
+        super(GO_MF, self).__init__()
+        self.data_path = data_path
+        self.user_prompt = prompt
+        self.return_prompt = return_prompt
+        
+        self.data_list = self._load_and_preprocess(self.data_path)
+        self.text2id = self._build_text_vocab()
+
+    def _load_and_preprocess(self, data_path):
+        data_list = []
+        df = pd.read_csv(data_path)
+        for _, row in df.iterrows():
+            try:
+                prot_seq = str(row['question']).strip()
+                result = str(row['answer']).strip()
+
+                text_seq = f"<answer>{result}</answer>\n"
+ 
+                prompt = f"""
+【Task】Predict the molecular functions performed by this protein.
+【Background】Each function is represented by a GO-MF term (e.g., GO:0003677), describing specific biochemical activities of the protein.
+【Output Format】List the predicted GO-MF terms, separated by commas, and wrap them in <answer> </answer> tags.  
+Example: <answer>GO:0003677, GO:0005524, GO:0016787</answer>
+"""
+                if self.user_prompt:
+                    prompt += self.user_prompt
+
+                # extra可以返回原始feather字符串，也可以返回feather_vals
+                 # 或 feather_raw
+                data_list.append((prot_seq, text_seq, prompt))
+            except Exception as e:
+                print(f"警告: 跳过有问题的行: {row}，原因: {e}")
+        return data_list
+
+    def _build_text_vocab(self):
+        text2id = {}
+        for _, text_seq, _ in self.data_list:
+            if text_seq not in text2id:
+                text2id[text_seq] = len(text2id)
+        return text2id
+
+    def shuffle(self):
+        random.shuffle(self.data_list)
+        return self
+
+    def __len__(self):
+        return len(self.data_list)
+
+    def __getitem__(self, index):
+        prot_seq, text_seq, prompt = self.data_list[index]
+        if self.return_prompt:
+            return prot_seq, prompt, text_seq,index
+        return prot_seq, text_seq, index
+
+
+class EC(Dataset):
+    def __init__(self, data_path, prompt='', return_prompt=False):
+        super(EC, self).__init__()
+        self.data_path = data_path
+        self.user_prompt = prompt
+        self.return_prompt = return_prompt
+        
+        self.data_list = self._load_and_preprocess(self.data_path)
+        self.text2id = self._build_text_vocab()
+
+    def _load_and_preprocess(self, data_path):
+        data_list = []
+        df = pd.read_csv(data_path)
+        for _, row in df.iterrows():
+            try:
+                name = str(row['name']).strip()
+                # 先按“-”分割，得到结构信息和 UniProt ID
+                structure_part, uniprot_id = name.split('-')  # '3r7t_A', 'Q9PMG4'
+
+                # 再按“_”分割结构信息，得到 PDB ID 和链ID
+                pdb_id, chain_id = structure_part.split('_')  # '3r7t', 'A'
+                prot_seq = str(row['aa_seq']).strip()
+                result = str(row['label']).strip()
+
+                text_seq = f"<answer>{result}</answer>\n"
+ 
+                prompt = f"""
+The information provided above is protein information, one of the chains of its crystal structure {pdb_id}, named {chain_id}, and numbered {uniprot_id} in the Uniprot sequence database.
+Based on this information, the possible enzyme activity is inferred and the corresponding EC number is predicted.
+【Output Format】List predicted EC numbers, separated by commas, wrapped in <answer> </answer> tags. 
+"""
+                if self.user_prompt:
+                    prompt += self.user_prompt
+
+                # extra可以返回原始feather字符串，也可以返回feather_vals
+                 # 或 feather_raw
+                data_list.append((prot_seq, text_seq, prompt))
+            except Exception as e:
+                print(f"警告: 跳过有问题的行: {row}，原因: {e}")
+        return data_list
+
+    def _build_text_vocab(self):
+        text2id = {}
+        for _, text_seq, _ in self.data_list:
+            if text_seq not in text2id:
+                text2id[text_seq] = len(text2id)
+        return text2id
+
+    def shuffle(self):
+        random.shuffle(self.data_list)
+        return self
+
+    def __len__(self):
+        return len(self.data_list)
+
+    def __getitem__(self, index):
+        prot_seq, text_seq, prompt = self.data_list[index]
+        if self.return_prompt:
+            return prot_seq, prompt, text_seq,index
+        return prot_seq, text_seq, index

data_provider/llm_tuning_dm.py

@@ -0,0 +1,261 @@
+# Copyright (c) Microsoft Corporation.
+# Licensed under the MIT License.
+from pytorch_lightning import LightningDataModule
+from data_provider.gal_helpers import escape_custom_split_sequence
+from data_provider.stage1_dm import SwissProtDataset, OntoProteinDataset
+from torch.utils.data import DataLoader, ConcatDataset
+
+
+class LLMTuningCollater:
+    def __init__(self, tokenizer, text_max_len, prot_max_len, use_gal):
+        self.text_max_len = text_max_len
+        self.prot_max_len = prot_max_len
+        self.tokenizer = tokenizer
+        self.use_gal = use_gal
+        
+    def __call__(self, batch):
+        prot_seqs, prompt_seqs, text_seqs, _ = zip(*batch)
+        prot_seqs = [prompt.format(p) for prompt, p in zip(prompt_seqs, prot_seqs)]
+        if self.use_gal:
+            prot_seqs = [escape_custom_split_sequence(p) for p in prot_seqs]
+        ## deal with prompt
+        self.tokenizer.padding_side = 'left'
+        prot_batch = self.tokenizer(text=prot_seqs,
+                                    truncation=True,
+                                    padding='max_length',
+                                    add_special_tokens=True,
+                                    max_length=self.prot_max_len,
+                                    return_tensors='pt',
+                                    return_attention_mask=True)
+        self.tokenizer.padding_side = 'right'
+        text_batch = self.tokenizer(text=text_seqs,
+                                    truncation=True,
+                                    padding='max_length',
+                                    add_special_tokens=True,
+                                    max_length=self.text_max_len,
+                                    return_tensors='pt',
+                                    return_attention_mask=True)
+        return prot_batch, text_batch
+
+
+class InferenceCollater:
+    def __init__(self, tokenizer, text_max_len, prot_max_len, use_gal):
+        self.text_max_len = text_max_len
+        self.prot_max_len = prot_max_len
+        self.tokenizer = tokenizer
+        self.use_gal = use_gal
+        
+    def __call__(self, batch):
+        prot_seqs, prompt_seqs, text_seqs, indices = zip(*batch)
+        prot_seqs = [prompt.format(p) for prompt, p in zip(prompt_seqs, prot_seqs)]
+        if self.use_gal:
+            prot_seqs = [escape_custom_split_sequence(p) for p in prot_seqs]
+        ## deal with prompt
+        self.tokenizer.padding_side = 'left'
+        prot_batch = self.tokenizer(text=prot_seqs,
+                                    truncation=True,
+                                    padding='max_length',
+                                    add_special_tokens=True,
+                                    max_length=self.prot_max_len,
+                                    return_tensors='pt',
+                                    return_attention_mask=True)
+        target_dict = {'targets': text_seqs, 'indices': indices}
+        return prot_batch, target_dict
+
+
+
+class LLMTuningDM(LightningDataModule):
+    def __init__(
+        self,
+        root: str = 'data/',
+        args=None,
+    ):
+        super().__init__()
+        self.batch_size = args.batch_size
+        self.inference_batch_size = args.inference_batch_size
+        self.num_workers = args.num_workers
+        self.prot_max_len = args.prot_max_len
+        self.text_max_len = args.text_max_len
+        if root.find('SwissProtV3') >= 0:
+            self.train_dataset = SwissProtDataset(root+'/train_set.json', prompt='[START_AMINO]{}[END_AMINO]. Swiss-Prot description: ', return_prompt=True)
+            self.val_dataset = SwissProtDataset(root+'/valid_set.json', prompt='[START_AMINO]{}[END_AMINO]. Swiss-Prot description: ', return_prompt=True)
+            self.test_dataset = SwissProtDataset(root+'/test_set.json', prompt='[START_AMINO]{}[END_AMINO]. Swiss-Prot description: ', return_prompt=True)
+        elif root.find('OntoProteinDatasetV2') >= 0:
+            self.train_dataset = OntoProteinDataset(root+'/train.txt', prompt='[START_AMINO]{}[END_AMINO]. Gene Ontology description: ', return_prompt=True)
+            self.val_dataset = OntoProteinDataset(root+'/valid.txt', prompt='[START_AMINO]{}[END_AMINO]. Gene Ontology description: ', return_prompt=True)
+            self.test_dataset = OntoProteinDataset(root+'/test.txt', prompt='[START_AMINO]{}[END_AMINO]. Gene Ontology description: ', return_prompt=True)
+        else:
+            raise NotImplementedError()
+        self.tokenizer = None
+        self.use_gal = args.llm_name.find('gal') >= 0
+    
+    def init_tokenizer(self, tokenizer):
+        self.tokenizer = tokenizer
+
+    def train_dataloader(self):
+        loader = DataLoader(
+            self.train_dataset,
+            batch_size=self.batch_size,
+            shuffle=True,
+            num_workers=self.num_workers,
+            pin_memory=False,
+            drop_last=True,
+            persistent_workers=False,
+            collate_fn=LLMTuningCollater(self.tokenizer, self.text_max_len, self.prot_max_len, self.use_gal),
+        )
+        return loader
+    
+    def val_dataloader(self):
+        val_loader = DataLoader(
+            self.val_dataset,
+            batch_size=self.batch_size,
+            shuffle=False,
+            num_workers=self.num_workers,
+            pin_memory=False,
+            drop_last=False,
+            persistent_workers=False,
+            collate_fn=LLMTuningCollater(self.tokenizer, self.text_max_len, self.prot_max_len, self.use_gal),
+        )
+        test_loader = DataLoader(
+            self.test_dataset,
+            batch_size=self.inference_batch_size,
+            shuffle=False,
+            num_workers=self.num_workers,
+            pin_memory=False,
+            drop_last=False,
+            persistent_workers=False,
+            collate_fn=InferenceCollater(self.tokenizer, self.text_max_len, self.prot_max_len, self.use_gal),
+        )
+        return [val_loader, test_loader]
+
+    def add_model_specific_args(parent_parser):
+        parser = parent_parser.add_argument_group("Data module")
+        parser.add_argument('--num_workers', type=int, default=2)
+        parser.add_argument('--batch_size', type=int, default=32)
+        parser.add_argument('--inference_batch_size', type=int, default=4)
+        parser.add_argument('--root', type=str, default='data/SwissProtV3')
+        parser.add_argument('--text_max_len', type=int, default=128)
+        parser.add_argument('--prot_max_len', type=int, default=1024)
+        parser.add_argument('--q_max_len', type=int, default=1064)
+        parser.add_argument('--a_max_len', type=int, default=36)
+        parser.add_argument('--prompt', type=str, default='[START_AMINO]{}[END_AMINO]. The protein has the following properties: ')
+        parser.add_argument('--filter_side_qa', action='store_true', default=False)
+        return parent_parser
+
+
+class LLMTuningMixDM(LightningDataModule):
+    def __init__(
+        self,
+        root: str = 'data/',
+        args=None,
+    ):
+        super().__init__()
+        self.batch_size = args.batch_size
+        self.inference_batch_size = args.inference_batch_size
+        self.num_workers = args.num_workers
+        self.prot_max_len = args.prot_max_len
+        self.text_max_len = args.text_max_len
+        
+        train_dataset1 = SwissProtDataset(root+'/SwissProtV3/train_set.json', prompt='[START_AMINO]{}[END_AMINO]. Swiss-Prot description: ', return_prompt=True)
+        train_dataset2 = OntoProteinDataset(root+'/OntoProteinDatasetV2/train.txt', prompt='[START_AMINO]{}[END_AMINO]. Gene Ontology description: ', return_prompt=True)
+        self.train_dataset = ConcatDataset([train_dataset1, train_dataset2])
+        self.swiss_val_dataset = SwissProtDataset(root+'/SwissProtV3/valid_set.json', prompt='[START_AMINO]{}[END_AMINO]. Swiss-Prot description: ', return_prompt=True)
+        self.onto_val_dataset = OntoProteinDataset(root+'/OntoProteinDatasetV2/valid.txt', prompt='[START_AMINO]{}[END_AMINO]. Gene Ontology description: ', return_prompt=True)
+        self.swiss_test_dataset = SwissProtDataset(root+'/SwissProtV3/test_set.json', prompt='[START_AMINO]{}[END_AMINO]. Swiss-Prot description: ', return_prompt=True)
+        self.onto_test_dataset = OntoProteinDataset(root+'/OntoProteinDatasetV2/test.txt', prompt='[START_AMINO]{}[END_AMINO]. Gene Ontology description: ', return_prompt=True)
+        
+        self.tokenizer = None
+        self.use_gal = args.llm_name.find('gal') >= 0
+    
+    def init_tokenizer(self, tokenizer):
+        self.tokenizer = tokenizer
+
+    def train_dataloader(self):
+        loader = DataLoader(
+            self.train_dataset,
+            batch_size=self.batch_size,
+            shuffle=True,
+            num_workers=self.num_workers,
+            pin_memory=False,
+            drop_last=True,
+            persistent_workers=False,
+            collate_fn=LLMTuningCollater(self.tokenizer, self.text_max_len, self.prot_max_len, self.use_gal),
+        )
+        return loader
+    
+    def val_dataloader(self):
+        swiss_val_loader = DataLoader(
+            self.swiss_val_dataset,
+            batch_size=self.batch_size,
+            shuffle=False,
+            num_workers=self.num_workers,
+            pin_memory=False,
+            drop_last=False,
+            persistent_workers=False,
+            collate_fn=LLMTuningCollater(self.tokenizer, self.text_max_len, self.prot_max_len, self.use_gal),
+        )
+        swiss_test_loader = DataLoader(
+            self.swiss_test_dataset,
+            batch_size=self.inference_batch_size,
+            shuffle=False,
+            num_workers=self.num_workers,
+            pin_memory=False,
+            drop_last=False,
+            persistent_workers=False,
+            collate_fn=InferenceCollater(self.tokenizer, self.text_max_len, self.prot_max_len, self.use_gal),
+        )
+
+        onto_val_loader = DataLoader(
+            self.onto_val_dataset,
+            batch_size=self.batch_size,
+            shuffle=False,
+            num_workers=self.num_workers,
+            pin_memory=False,
+            drop_last=False,
+            persistent_workers=False,
+            collate_fn=LLMTuningCollater(self.tokenizer, self.text_max_len, self.prot_max_len, self.use_gal),
+        )
+        onto_test_loader = DataLoader(
+            self.onto_test_dataset,
+            batch_size=self.inference_batch_size,
+            shuffle=False,
+            num_workers=self.num_workers,
+            pin_memory=False,
+            drop_last=False,
+            persistent_workers=False,
+            collate_fn=InferenceCollater(self.tokenizer, self.text_max_len, self.prot_max_len, self.use_gal),
+        )
+        return [swiss_val_loader, swiss_test_loader, onto_val_loader, onto_test_loader]
+
+    def add_model_specific_args(parent_parser):
+        parser = parent_parser.add_argument_group("Data module")
+        parser.add_argument('--num_workers', type=int, default=2)
+        parser.add_argument('--batch_size', type=int, default=32)
+        parser.add_argument('--inference_batch_size', type=int, default=4)
+        parser.add_argument('--root', type=str, default='data/SwissProtV3')
+        parser.add_argument('--text_max_len', type=int, default=128)
+        parser.add_argument('--prot_max_len', type=int, default=1024)
+        parser.add_argument('--q_max_len', type=int, default=1064)
+        parser.add_argument('--a_max_len', type=int, default=36)
+        parser.add_argument('--prompt', type=str, default='[START_AMINO]{}[END_AMINO]. The protein has the following properties: ')
+        parser.add_argument('--filter_side_qa', action='store_true', default=False)
+        return parent_parser
+    
+
+if __name__ == '__main__':
+    dataset = SwissProtDataset('../data/SwissProtV3/train_set.json')
+    from transformers import AutoTokenizer
+    tokenizer = AutoTokenizer.from_pretrained('facebook/galactica-1.3b')
+    tokenizer.add_special_tokens({'pad_token': '<pad>'})
+    loader = DataLoader(
+            dataset,
+            batch_size=16,
+            shuffle=True,
+            num_workers=0,
+            pin_memory=False,
+            drop_last=True,
+            persistent_workers=False,
+            collate_fn=LLMTuningCollater(tokenizer, 128, 1024, True, '[START_AMINO]{}[END_AMINO].'),
+        )
+    for data in loader:
+        input()

data_provider/llm_tuning_prot_qa_dm.py

@@ -0,0 +1,164 @@
+# Copyright (c) Microsoft Corporation.
+# Licensed under the MIT License.
+from pytorch_lightning import LightningDataModule
+from torch.utils.data import DataLoader
+from data_provider.prot_qa_dm import PDBQADataset
+from data_provider.gal_helpers import escape_custom_split_sequence
+
+
+class LLMTuningProtQACollater(object):
+    def __init__(self, tokenizer, q_max_len, a_max_len, use_gal, prompt):
+        self.tokenizer = tokenizer
+        self.q_max_len = q_max_len
+        self.a_max_len = a_max_len
+        self.use_gal = use_gal
+        self.prompt = prompt
+        assert prompt.find('{}') >= 0
+        
+    def __call__(self, batch):
+        prot_seqs, questions, answers, q_types = zip(*batch)
+        assert len(prot_seqs) == len(questions) == len(answers)
+        questions = [self.prompt.format(prot_seqs[i], questions[i]) for i in range(len(prot_seqs))]
+
+        if self.use_gal:
+            questions = [escape_custom_split_sequence(q) for q in questions]
+        answers = [a + '\n' for a in answers]
+        if False:
+            self.tokenizer.padding_side = 'left'
+            q_batch = self.tokenizer(questions,
+                                    truncation=True,
+                                    padding='max_length',
+                                    add_special_tokens=True,
+                                    max_length=self.q_max_len,
+                                    return_tensors='pt',
+                                    return_attention_mask=True, 
+                                    return_token_type_ids=False)
+            self.tokenizer.padding_side = 'right'
+            a_batch = self.tokenizer(answers,
+                                    truncation=True,
+                                    padding='max_length',
+                                    add_special_tokens=True,
+                                    max_length=self.a_max_len,
+                                    return_tensors='pt',
+                                    return_attention_mask=True, 
+                                    return_token_type_ids=False)
+            return q_batch, a_batch
+        else:
+            self.tokenizer.padding_side = 'right'
+            qa_pair = [[q, a] for q, a in zip(questions, answers)]
+            qa_batch = self.tokenizer(qa_pair,
+                                      truncation=True,
+                                      padding='max_length',
+                                      add_special_tokens=True,
+                                      max_length=self.q_max_len + self.a_max_len,
+                                      return_tensors='pt',
+                                      return_attention_mask=True,
+                                      return_token_type_ids=True)
+            return qa_batch
+
+
+class InferenceCollater(object):
+    def __init__(self, tokenizer, q_max_len, a_max_len, use_gal, prompt):
+        self.tokenizer = tokenizer
+        self.q_max_len = q_max_len
+        self.a_max_len = a_max_len
+        self.use_gal = use_gal
+        self.prompt = prompt
+        assert prompt.find('{}') >= 0
+        
+    def __call__(self, batch):
+        prot_seqs, questions, answers, q_types, indices = zip(*batch)
+        assert len(prot_seqs) == len(questions) == len(answers)
+        questions = [self.prompt.format(prot_seqs[i], questions[i]) for i in range(len(prot_seqs))]
+
+        if self.use_gal:
+            questions = [escape_custom_split_sequence(q) for q in questions]
+        answers = [a + '\n' for a in answers]
+        self.tokenizer.padding_side = 'left'
+        q_batch = self.tokenizer(questions,
+                                 truncation=True,
+                                 padding='max_length',
+                                 add_special_tokens=True,
+                                 max_length=self.q_max_len,
+                                 return_tensors='pt',
+                                 return_attention_mask=True, 
+                                 return_token_type_ids=False)
+        target_dict = {'targets': answers, 'q_types': q_types, 'indices': indices}
+        return q_batch, target_dict
+
+
+class LLMTuningProtQADM(LightningDataModule):
+    def __init__(
+        self,
+        root: str = 'data/',
+        args=None,
+    ):
+        super().__init__()
+        self.args = args
+        self.batch_size = args.batch_size
+        self.inference_batch_size = args.inference_batch_size
+        self.num_workers = args.num_workers
+        self.q_max_len = args.q_max_len
+        self.a_max_len = args.a_max_len
+        self.prompt = args.prompt
+        
+        self.train_dataset = PDBQADataset(root, 'train.txt', "Question: {} Answer:", filter_side_qa=args.filter_side_qa)
+        self.val_dataset = PDBQADataset(root, 'val.txt', "Question: {} Answer:", filter_side_qa=args.filter_side_qa)
+        self.test_dataset = PDBQADataset(root, 'test.txt', "Question: {} Answer:", filter_side_qa=args.filter_side_qa)
+        
+        self.tokenizer = None
+        self.use_gal = args.llm_name.find('gal') >= 0
+    
+    def init_tokenizer(self, tokenizer):
+        self.tokenizer = tokenizer
+
+    def train_dataloader(self):
+        loader = DataLoader(
+            self.train_dataset,
+            batch_size=self.batch_size,
+            shuffle=True,
+            num_workers=self.num_workers,
+            pin_memory=False,
+            drop_last=True,
+            persistent_workers=False,
+            collate_fn=LLMTuningProtQACollater(self.tokenizer, self.q_max_len, self.a_max_len, self.use_gal, self.prompt),
+        )
+        return loader
+
+    def val_dataloader(self):
+        val_loader = DataLoader(
+            self.val_dataset,
+            batch_size=self.batch_size,
+            shuffle=False,
+            num_workers=self.num_workers,
+            pin_memory=False,
+            drop_last=False,
+            persistent_workers=False,
+            collate_fn=LLMTuningProtQACollater(self.tokenizer, self.q_max_len, self.a_max_len, self.use_gal, self.prompt),
+        )
+        test_loader = DataLoader(
+            self.test_dataset,
+            batch_size=self.inference_batch_size,
+            shuffle=False,
+            num_workers=self.num_workers,
+            pin_memory=False,
+            drop_last=False,
+            persistent_workers=False,
+            collate_fn=InferenceCollater(self.tokenizer, self.q_max_len, self.a_max_len, self.use_gal, self.prompt),
+        )
+        return [val_loader, test_loader]
+    
+
+    def add_model_specific_args(parent_parser):
+        parser = parent_parser.add_argument_group("Data module")
+        parser.add_argument('--num_workers', type=int, default=2)
+        parser.add_argument('--batch_size', type=int, default=32)
+        parser.add_argument('--inference_batch_size', type=int, default=4)
+        parser.add_argument('--root', type=str, default='data/SwissProtV3')
+        parser.add_argument('--q_max_len', type=int, default=1064)
+        parser.add_argument('--a_max_len', type=int, default=36)
+        parser.add_argument('--prompt', type=str, default='[START_AMINO]{}[END_AMINO]. {}')
+        parser.add_argument('--filter_side_qa', action='store_true', default=False)
+        return parent_parser
+
+

data_provider/metalIonbinding.py

@@ -0,0 +1,63 @@
+import json
+from pytorch_lightning import LightningDataModule
+from torch.utils.data import DataLoader, ConcatDataset, Dataset
+from data_provider.stage1_dm import SwissProtDataset, OntoProteinDataset
+import pandas as pd
+
+class MetallonBinding(Dataset):
+    def __init__(self, data_path, prompt='', return_prompt=False):
+        super(MetallonBinding, self).__init__()
+        self.data_path = data_path
+        self.user_prompt = prompt
+        self.return_prompt = return_prompt
+        
+        self.data_list = self._load_and_preprocess(self.data_path)
+        self.text2id = self._build_text_vocab()
+
+    def _load_and_preprocess(self, data_path):
+        data_list = []
+        df = pd.read_csv(data_path)
+        for _, row in df.iterrows():
+            try:
+                name = str(row['name']).strip()
+                prot_seq = str(row['aa_seq']).strip()
+                result = str(int(row['label'])).strip()
+
+                text_seq = f"<answer>{result}</answer>\n"
+ 
+                prompt = f"""
+Task: Determine whether this protein is a metalloprotein based on the provided sequence and protein name {name}.
+Background: Metalloproteins are proteins that bind metal ions, often through specific amino acid residues such as histidine (H), cysteine (C), aspartate (D), or glutamate (E).
+Question: Does this protein bind metal ions? Please choose one of the following options:  
+0: Non-metalloprotein — This protein does **not** bind to any metal ions.   
+1: Metalloprotein — This protein **binds** to one or more metal ions.
+                """
+                if self.user_prompt:
+                    prompt += self.user_prompt
+
+                # extra可以返回原始feather字符串，也可以返回feather_vals
+                 # 或 feather_raw
+                data_list.append((prot_seq, text_seq, prompt))
+            except Exception as e:
+                print(f"警告: 跳过有问题的行: {row}，原因: {e}")
+        return data_list
+
+    def _build_text_vocab(self):
+        text2id = {}
+        for _, text_seq, _ in self.data_list:
+            if text_seq not in text2id:
+                text2id[text_seq] = len(text2id)
+        return text2id
+
+    def shuffle(self):
+        random.shuffle(self.data_list)
+        return self
+
+    def __len__(self):
+        return len(self.data_list)
+
+    def __getitem__(self, index):
+        prot_seq, text_seq, prompt = self.data_list[index]
+        if self.return_prompt:
+            return prot_seq, prompt, text_seq,index
+        return prot_seq, text_seq, index

data_provider/mutation.py

@@ -0,0 +1,119 @@
+import json
+from pytorch_lightning import LightningDataModule
+from torch.utils.data import DataLoader, ConcatDataset, Dataset
+from data_provider.stage1_dm import SwissProtDataset, OntoProteinDataset
+import pandas as pd
+
+class TAPE_Stability(Dataset):
+    def __init__(self, data_path, prompt='', return_prompt=False):
+        super(TAPE_Stability, self).__init__()
+        self.data_path = data_path
+        self.user_prompt = prompt
+        self.return_prompt = return_prompt
+        
+        self.data_list = self._load_and_preprocess(self.data_path)
+        self.text2id = self._build_text_vocab()
+
+    def _load_and_preprocess(self, data_path):
+        data_list = []
+        df = pd.read_csv(data_path)
+        for _, row in df.iterrows():
+            try:
+                
+                prot_seq = str(row['aa_seq']).strip()
+                result = str(row['label']).strip()
+
+                text_seq = f"<answer>{result}</answer>\n"
+ 
+                prompt = """
+【Task】Predict the thermostability score of the given protein sequence, which reflects its ability to maintain proper folding above a concentration threshold.
+【Background】Protein stability is an important biophysical property indicating a protein’s resistance to denaturation or unfolding under thermal or chemical stress. In this task, each protein is evaluated by a numerical stability score, where higher values indicate greater ability to remain folded under extreme conditions. This score serves as a proxy for the protein’s intrinsic stability.
+【Question】What is the predicted stability score for this sequence?
+【Output Format】You must return only the score number, wrapped in <answer></answer> tags.  
+"""
+                if self.user_prompt:
+                    prompt += self.user_prompt
+
+                # extra可以返回原始feather字符串，也可以返回feather_vals
+                 # 或 feather_raw
+                data_list.append((prot_seq, text_seq, prompt))
+            except Exception as e:
+                print(f"警告: 跳过有问题的行: {row}，原因: {e}")
+        return data_list
+
+    def _build_text_vocab(self):
+        text2id = {}
+        for _, text_seq, _ in self.data_list:
+            if text_seq not in text2id:
+                text2id[text_seq] = len(text2id)
+        return text2id
+
+    def shuffle(self):
+        random.shuffle(self.data_list)
+        return self
+
+    def __len__(self):
+        return len(self.data_list)
+
+    def __getitem__(self, index):
+        prot_seq, text_seq, prompt = self.data_list[index]
+        if self.return_prompt:
+            return prot_seq, prompt, text_seq,index
+        return prot_seq, text_seq, index
+
+class TAPE_Fluorescence(Dataset):
+    def __init__(self, data_path, prompt='', return_prompt=False):
+        super(TAPE_Fluorescence, self).__init__()
+        self.data_path = data_path
+        self.user_prompt = prompt
+        self.return_prompt = return_prompt
+        
+        self.data_list = self._load_and_preprocess(self.data_path)
+        self.text2id = self._build_text_vocab()
+
+    def _load_and_preprocess(self, data_path):
+        data_list = []
+        df = pd.read_csv(data_path)
+        for _, row in df.iterrows():
+            try:
+                
+                prot_seq = str(row['aa_seq']).strip()
+                result = str(row['label']).strip()
+
+                text_seq = f"<answer>{result}</answer>\n"
+ 
+                prompt = """
+【Task】Predict the log fluorescence intensity of the given protein sequence.
+【Output Format】You must return only the numerical value, wrapped in <answer></answer> tags.  
+"""
+# 【Background】Fluorescence intensity reflects how strongly a protein emits light when excited by a specific wavelength. It is commonly measured in protein variants such as GFP (Green Fluorescent Protein) mutants. The log-transformed fluorescence value quantifies the brightness on a logarithmic scale. Mutations in the sequence can increase or decrease fluorescence intensity.
+# 【Question】What is the predicted log fluorescence intensity for this sequence?
+                if self.user_prompt:
+                    prompt += self.user_prompt
+
+                # extra可以返回原始feather字符串，也可以返回feather_vals
+                 # 或 feather_raw
+                data_list.append((prot_seq, text_seq, prompt))
+            except Exception as e:
+                print(f"警告: 跳过有问题的行: {row}，原因: {e}")
+        return data_list
+
+    def _build_text_vocab(self):
+        text2id = {}
+        for _, text_seq, _ in self.data_list:
+            if text_seq not in text2id:
+                text2id[text_seq] = len(text2id)
+        return text2id
+
+    def shuffle(self):
+        random.shuffle(self.data_list)
+        return self
+
+    def __len__(self):
+        return len(self.data_list)
+
+    def __getitem__(self, index):
+        prot_seq, text_seq, prompt = self.data_list[index]
+        if self.return_prompt:
+            return prot_seq, prompt, text_seq,index
+        return prot_seq, text_seq, index

data_provider/production.py

@@ -0,0 +1,237 @@
+
+import json
+from pytorch_lightning import LightningDataModule
+from torch.utils.data import DataLoader, ConcatDataset, Dataset
+from data_provider.stage1_dm import SwissProtDataset, OntoProteinDataset
+import pandas as pd
+class Antibiotic_Resistance(Dataset):
+    def __init__(self, data_path, prompt='', return_prompt=False):
+        super(Antibiotic_Resistance, self).__init__()
+        self.data_path = data_path
+        self.user_prompt = prompt
+        self.return_prompt = return_prompt
+        
+        self.data_list = self._load_and_preprocess(self.data_path)
+        self.text2id = self._build_text_vocab()
+
+    def _load_and_preprocess(self, data_path):
+        data_list = []
+        df = pd.read_csv(data_path)
+        for _, row in df.iterrows():
+            try:
+                prot_seq = str(row['aa_seq']).strip()
+                result = str(row['label']).strip()
+                text_seq = f"<answer>{result}</answer>\n"
+ 
+                prompt = f"""
+【Task】Predict the antibiotic resistance class of the given protein.
+【Background】Antibiotic resistance refers to the ability of bacteria or other microbes to resist the effects of antibiotics that were once effective against them. Each protein is associated with resistance to exactly one of 19 antibiotic classes.
+
+【Prediction Goal】Based on the provided protein sequence, determine which single antibiotic class (from 1 to 19) this protein confers resistance to.
+
+【Output Format】Return only one predicted resistance class (a number from 1 to 19), wrapped in <answer> </answer> tags.  
+Example: <answer>7</answer>
+"""
+                if self.user_prompt:
+                    prompt += self.user_prompt
+
+                # extra可以返回原始feather字符串，也可以返回feather_vals
+                 # 或 feather_raw
+                data_list.append((prot_seq, text_seq, prompt))
+            except Exception as e:
+                print(f"警告: 跳过有问题的行: {row}，原因: {e}")
+        return data_list
+
+    def _build_text_vocab(self):
+        text2id = {}
+        for _, text_seq, _ in self.data_list:
+            if text_seq not in text2id:
+                text2id[text_seq] = len(text2id)
+        return text2id
+
+    def shuffle(self):
+        random.shuffle(self.data_list)
+        return self
+
+    def __len__(self):
+        return len(self.data_list)
+
+    def __getitem__(self, index):
+        prot_seq, text_seq, prompt = self.data_list[index]
+        if self.return_prompt:
+            return prot_seq, prompt, text_seq,index
+        return prot_seq, text_seq, index
+
+class Thermostability(Dataset):
+    def __init__(self, data_path, prompt='', return_prompt=False):
+        super(Thermostability, self).__init__()
+        self.data_path = data_path
+        self.user_prompt = prompt
+        self.return_prompt = return_prompt
+        
+        self.data_list = self._load_and_preprocess(self.data_path)
+        self.text2id = self._build_text_vocab()
+
+    def _load_and_preprocess(self, data_path):
+        data_list = []
+        df = pd.read_csv(data_path)
+        for _, row in df.iterrows():
+            try:
+                #name = str(row['name']).strip()
+                prot_seq = str(row['aa_seq']).strip()
+                result = str(row['label']).strip()
+                text_seq = f"<answer>{result}</answer>\n"
+ 
+                prompt = f"""
+【Task】Predict the thermostability value of the given protein.
+【Background】Thermostability refers to the ability of a molecule to resist irreversible chemical or physical changes at high temperatures, such as decomposition or aggregation.
+【Output Format】Provide the predicted thermostability as a numeric value (e.g., melting temperature in °C). Wrap your answer in <answer></answer> tags.  
+
+"""
+                if self.user_prompt:
+                    prompt += self.user_prompt
+
+                # extra可以返回原始feather字符串，也可以返回feather_vals
+                 # 或 feather_raw
+                data_list.append((prot_seq, text_seq, prompt))
+            except Exception as e:
+                print(f"警告: 跳过有问题的行: {row}，原因: {e}")
+        return data_list
+
+    def _build_text_vocab(self):
+        text2id = {}
+        for _, text_seq, _ in self.data_list:
+            if text_seq not in text2id:
+                text2id[text_seq] = len(text2id)
+        return text2id
+
+    def shuffle(self):
+        random.shuffle(self.data_list)
+        return self
+
+    def __len__(self):
+        return len(self.data_list)
+
+    def __getitem__(self, index):
+        prot_seq, text_seq, prompt = self.data_list[index]
+        if self.return_prompt:
+            return prot_seq, prompt, text_seq,index
+        return prot_seq, text_seq, index
+
+
+
+class Material(Dataset):
+    def __init__(self, data_path, prompt='', return_prompt=False):
+        super(Material, self).__init__()
+        self.data_path = data_path
+        self.user_prompt = prompt
+        self.return_prompt = return_prompt
+        
+        self.data_list = self._load_and_preprocess(self.data_path)
+        self.text2id = self._build_text_vocab()
+
+    def _load_and_preprocess(self, data_path):
+        data_list = []
+        df = pd.read_csv(data_path)
+        for _, row in df.iterrows():
+            try:
+                #name = str(row['name']).strip()
+                prot_seq = str(row['aa_seq']).strip()
+                result = str(row['label']).strip()
+                text_seq = f"<answer>{result}</answer>\n"
+ 
+                prompt = f"""
+【Task】Determine whether the given material can be successfully produced.
+【Background】In materials science, certain chemical compounds or materials may or may not be synthesizable (i.e., producible) under realistic experimental conditions. This task requires classifying whether the input material composition and structure allow for successful production. This is a binary classification problem.
+【Question】Can this material be successfully produced?
+【Output Format】Respond with either "1" or "0", and wrap your answer in <answer></answer> tags.  
+
+"""
+                if self.user_prompt:
+                    prompt += self.user_prompt
+
+                # extra可以返回原始feather字符串，也可以返回feather_vals
+                 # 或 feather_raw
+                data_list.append((prot_seq, text_seq, prompt))
+            except Exception as e:
+                print(f"警告: 跳过有问题的行: {row}，原因: {e}")
+        return data_list
+
+    def _build_text_vocab(self):
+        text2id = {}
+        for _, text_seq, _ in self.data_list:
+            if text_seq not in text2id:
+                text2id[text_seq] = len(text2id)
+        return text2id
+
+    def shuffle(self):
+        random.shuffle(self.data_list)
+        return self
+
+    def __len__(self):
+        return len(self.data_list)
+
+    def __getitem__(self, index):
+        prot_seq, text_seq, prompt = self.data_list[index]
+        if self.return_prompt:
+            return prot_seq, prompt, text_seq,index
+        return prot_seq, text_seq, index
+
+
+
+class Clone(Dataset):
+    def __init__(self, data_path, prompt='', return_prompt=False):
+        super(Clone, self).__init__()
+        self.data_path = data_path
+        self.user_prompt = prompt
+        self.return_prompt = return_prompt
+        
+        self.data_list = self._load_and_preprocess(self.data_path)
+        self.text2id = self._build_text_vocab()
+
+    def _load_and_preprocess(self, data_path):
+        data_list = []
+        df = pd.read_csv(data_path)
+        for _, row in df.iterrows():
+            try:
+                #name = str(row['name']).strip()
+                prot_seq = str(row['aa_seq']).strip()
+                result = str(row['label']).strip()
+                text_seq = f"<answer>{result}</answer>\n"
+ 
+                prompt = f"""
+
+【Task】Determine whether the given protein sequence can be successfully cloned.
+【Background】In molecular biology, cloning refers to the process of creating copies of a DNA or protein sequence. Some sequences can be challenging to clone due to their length, GC-content, secondary structures, or toxicity to the host. This task requires predicting whether the given protein sequence is likely to be successfully cloned. This is a binary classification problem.
+【Question】Can this protein sequence be successfully cloned?
+【Output Format】Respond with either "1" or "0", and wrap your answer in <answer></answer> tags.  
+"""
+                if self.user_prompt:
+                    prompt += self.user_prompt
+
+                # extra可以返回原始feather字符串，也可以返回feather_vals
+                 # 或 feather_raw
+                data_list.append((prot_seq, text_seq, prompt))
+            except Exception as e:
+                print(f"警告: 跳过有问题的行: {row}，原因: {e}")
+        return data_list
+
+    def _build_text_vocab(self):
+        text2id = {}
+        for _, text_seq, _ in self.data_list:
+            if text_seq not in text2id:
+                text2id[text_seq] = len(text2id)
+        return text2id
+
+    def shuffle(self):
+        random.shuffle(self.data_list)
+        return self
+
+    def __len__(self):
+        return len(self.data_list)
+
+    def __getitem__(self, index):
+        prot_seq, text_seq, prompt = self.data_list[index]
+        if self.return_prompt:
+            return prot_seq, prompt, text_seq,index
+        return prot_seq, text_seq, index

data_provider/prot_qa_dm.py

@@ -0,0 +1,299 @@
+# Copyright (c) Microsoft Corporation.
+# Licensed under the MIT License.
+import json
+from pytorch_lightning import LightningDataModule
+# import torch_geometric
+from torch.utils.data import DataLoader, Dataset
+from pathlib import Path
+
+
+class ProtQACollater(object):
+    def __init__(self, tokenizer, prot_tokenizer, q_max_len, a_max_len, prot_max_len):
+        self.tokenizer = tokenizer
+        self.prot_tokenizer = prot_tokenizer
+        self.q_max_len = q_max_len
+        self.a_max_len = a_max_len
+        self.prot_max_len = prot_max_len
+        
+    def __call__(self, batch):
+        prot_seqs, questions, answers, _, _ = zip(*batch)
+        answers = [a + '\n' for a in answers]
+        prot_batch = self.prot_tokenizer(prot_seqs,
+                                         truncation=True,
+                                         padding='max_length',
+                                         max_length=self.prot_max_len,
+                                         return_tensors="pt",
+                                         return_attention_mask=True, 
+                                         return_token_type_ids=False)
+        if False:
+            self.tokenizer.padding_side = 'left'
+            q_batch = self.tokenizer(questions,
+                                    truncation=True,
+                                    padding='max_length',
+                                    add_special_tokens=True,
+                                    max_length=self.q_max_len,
+                                    return_tensors='pt',
+                                    return_attention_mask=True, 
+                                    return_token_type_ids=False)
+            self.tokenizer.padding_side = 'right'
+            a_batch = self.tokenizer(answers,
+                                    truncation=True,
+                                    padding='max_length',
+                                    add_special_tokens=True,
+                                    max_length=self.a_max_len,
+                                    return_tensors='pt',
+                                    return_attention_mask=True, 
+                                    return_token_type_ids=False)
+            return prot_batch, q_batch, a_batch
+        else:
+            self.tokenizer.padding_side = 'right'
+            qa_pair = [[q, a] for q, a in zip(questions, answers)]
+            qa_batch = self.tokenizer(qa_pair,
+                                      truncation=True,
+                                      padding='max_length',
+                                      add_special_tokens=True,
+                                      max_length=self.q_max_len + self.a_max_len,
+                                      return_tensors='pt',
+                                      return_attention_mask=True,
+                                      return_token_type_ids=True)
+            return prot_batch, qa_batch
+
+
+class InferenceCollater(object):
+    def __init__(self, tokenizer, prot_tokenizer, q_max_len, a_max_len, prot_max_len):
+        self.tokenizer = tokenizer
+        self.prot_tokenizer = prot_tokenizer
+        self.q_max_len = q_max_len
+        self.a_max_len = a_max_len
+        self.prot_max_len = prot_max_len
+        
+    def __call__(self, batch):
+        prot_seqs, questions, answers, q_types, indices = zip(*batch)
+        answers = [a + '\n' for a in answers]
+        prot_batch = self.prot_tokenizer(prot_seqs,
+                                         truncation=True,
+                                         padding='max_length',
+                                         max_length=self.prot_max_len,
+                                         return_tensors="pt",
+                                         return_attention_mask=True, 
+                                         return_token_type_ids=False)
+        self.tokenizer.padding_side = 'left'
+        q_batch = self.tokenizer(questions,
+                                 truncation=True,
+                                 padding='max_length',
+                                 add_special_tokens=True,
+                                 max_length=self.q_max_len,
+                                 return_tensors='pt',
+                                 return_attention_mask=True, 
+                                 return_token_type_ids=False)        
+        target_dict = {'targets': answers, 'q_types': q_types, 'indices': indices}
+        return prot_batch, q_batch, target_dict
+
+
+class ProtQADM(LightningDataModule):
+    def __init__(
+        self,
+        root: str = 'data/',
+        args=None,
+    ):
+        super().__init__()
+        self.args = args
+        self.batch_size = args.batch_size
+        self.inference_batch_size = args.inference_batch_size
+        self.num_workers = args.num_workers
+        self.q_max_len = args.q_max_len
+        self.a_max_len = args.a_max_len
+        self.prot_max_len = args.prot_max_len
+        self.prompt = args.prompt
+        
+        self.train_dataset = PDBQADataset(root, 'train.txt', prompt=self.prompt, filter_side_qa=args.filter_side_qa)
+        self.val_dataset = PDBQADataset(root, 'val.txt', prompt=self.prompt, filter_side_qa=args.filter_side_qa)
+        self.test_dataset = PDBQADataset(root, 'test.txt', prompt=self.prompt, filter_side_qa=args.filter_side_qa)
+        
+        self.tokenizer = None
+        self.prot_tokenizer = None
+    
+    def init_tokenizer(self, tokenizer, prot_tokenizer):
+        self.tokenizer = tokenizer
+        self.prot_tokenizer = prot_tokenizer
+
+    def train_dataloader(self):
+        loader = DataLoader(
+            self.train_dataset,
+            batch_size=self.batch_size,
+            shuffle=True,
+            num_workers=self.num_workers,
+            pin_memory=False,
+            drop_last=True,
+            persistent_workers=False,
+            collate_fn=ProtQACollater(self.tokenizer, self.prot_tokenizer, self.q_max_len, self.a_max_len, self.prot_max_len),
+        )
+        return loader
+
+    def val_dataloader(self):
+        val_loader = DataLoader(
+            self.val_dataset,
+            batch_size=self.batch_size,
+            shuffle=False,
+            num_workers=self.num_workers,
+            pin_memory=False,
+            drop_last=False,
+            persistent_workers=False,
+            collate_fn=ProtQACollater(self.tokenizer, self.prot_tokenizer, self.q_max_len, self.a_max_len, self.prot_max_len),
+        )
+        test_loader = DataLoader(
+            self.test_dataset,
+            batch_size=self.inference_batch_size,
+            shuffle=False,
+            num_workers=self.num_workers,
+            pin_memory=False,
+            drop_last=False,
+            persistent_workers=False,
+            collate_fn=InferenceCollater(self.tokenizer, self.prot_tokenizer, self.q_max_len, self.a_max_len, self.prot_max_len),
+        )
+        return [val_loader, test_loader]
+    
+
+    def add_model_specific_args(parent_parser):
+        parser = parent_parser.add_argument_group("Data module")
+        parser.add_argument('--num_workers', type=int, default=2)
+        parser.add_argument('--batch_size', type=int, default=32)
+        parser.add_argument('--inference_batch_size', type=int, default=4)
+        parser.add_argument('--root', type=str, default='data/SwissProtV3')
+        parser.add_argument('--text_max_len', type=int, default=128)
+        parser.add_argument('--q_max_len', type=int, default=34)
+        parser.add_argument('--a_max_len', type=int, default=36)
+        parser.add_argument('--prot_max_len', type=int, default=1024)
+        parser.add_argument('--prompt', type=str, default='The protein has the following properties: ')
+        parser.add_argument('--filter_side_qa', action='store_true', default=False)
+        return parent_parser
+
+
+class PDBQADataset(Dataset):
+    def __init__(self, root_path, subset, prompt="Question: {} Answer:", filter_side_qa=False):
+        super(PDBQADataset, self).__init__()
+        self.data_path = Path(root_path) / subset
+        self.qa_path = Path(root_path) / 'qa_all.json'
+        self.q_type_path = Path(root_path) / 'q_types.txt'
+        self.prompt = prompt
+
+        ## load dataset
+        with open(self.qa_path, 'r') as f:
+            qa_data = json.load(f)
+        
+        with open(self.data_path, 'r') as f:
+            lines = f.readlines()
+            pdb2seq = [line.strip().split('\t') for line in lines]
+        
+
+        ## load q types
+        with open(self.q_type_path, 'r') as f:
+            q_types = [line.strip().split('\t') for line in f.readlines()]
+        self.q_type_dict = {q: t for q, t in q_types}
+
+        ## process dataset
+        pdb_set = set(i[0] for i in pdb2seq)
+        ## filter qa data
+        qa_data = {k: v for k, v in qa_data.items() if k in pdb_set}
+        assert len(qa_data) == len(pdb_set), print(len(qa_data), len(pdb_set))
+        
+        ## generate qa data
+        self.data_list = []
+        for pdb_id, seq in pdb2seq:
+            qa_list = qa_data[pdb_id]
+            for qa in qa_list:
+                q = qa['Q']
+                a = str(qa['A'])
+                if filter_side_qa:
+                    q_type = self.q_type_dict[q]
+                    if q_type.find('side information') >= 0:
+                        continue
+                self.data_list.append((seq, q, a))
+
+    def __len__(self):
+        return len(self.data_list)
+
+    def __getitem__(self, index):
+        seq, q, a = self.data_list[index]
+        q_type = self.q_type_dict[q]
+        q = self.prompt.format(q)
+        return seq, q, a, q_type, index
+
+
+if __name__ == '__main__':
+    import numpy as np
+    from collections import defaultdict, Counter
+    train_dataset = PDBQADataset('../data/PDBDataset', 'train.txt', filter_side_qa=True)
+    val_dataset = PDBQADataset('../data/PDBDataset', 'val.txt', filter_side_qa=True)
+    test_dataset = PDBQADataset('../data/PDBDataset', 'test.txt', filter_side_qa=True)
+    if True:
+        # print(len(train_dataset), len(val_dataset), len(test_dataset))
+        # train_protein_lens = np.asarray([len(p) for p in train_dataset.protein_list])
+        # val_protein_lens = np.asarray([len(p) for p in val_dataset.protein_list])
+        # test_protein_lens = np.asarray([len(p) for p in test_dataset.protein_list])
+        
+        q_lens =  []
+        a_lens = []
+        for seq, q, a in train_dataset.data_list:
+            q_lens.append(len(q.split()))
+            a_lens.append(len(a.split()))
+
+        print(np.asarray(q_lens).min(), np.asarray(q_lens).max(), np.asarray(q_lens).mean())
+        print(np.asarray(a_lens).min(), np.asarray(a_lens).max(), np.asarray(a_lens).mean())
+        
+        q_lens =  []
+        a_lens = []
+        for seq, q, a in val_dataset.data_list:
+            q_lens.append(len(q.split()))
+            a_lens.append(len(a.split()))
+
+        print(np.asarray(q_lens).min(), np.asarray(q_lens).max(), np.asarray(q_lens).mean())
+        print(np.asarray(a_lens).min(), np.asarray(a_lens).max(), np.asarray(a_lens).mean())
+
+        q_lens =  []
+        a_lens = []
+        for seq, q, a in test_dataset.data_list:
+            q_lens.append(len(q.split()))
+            a_lens.append(len(a.split()))
+
+        print(np.asarray(q_lens).min(), np.asarray(q_lens).max(), np.asarray(q_lens).mean())
+        print(np.asarray(a_lens).min(), np.asarray(a_lens).max(), np.asarray(a_lens).mean())
+
+    elif False:
+        ## construct the guess for prediction by number
+        train_counter = defaultdict(Counter)
+        for  _, q, a in train_dataset.data_list:
+            train_counter[q.lower()][a] += 1
+        ## get the most common answer
+        q2a = {}
+        for q, counter in train_counter.items():
+            q2a[q] = counter.most_common(1)[0][0]
+
+        ## test the guess
+        acc = 0
+        for _, q, a in test_dataset.data_list:
+            if q.lower() in q2a:
+                predict = q2a[q.lower()]
+                if predict.lower() == a.lower():
+                    acc += 1
+        print(acc / len(test_dataset.data_list))
+    elif False:
+        from transformers import AutoTokenizer, EsmTokenizer
+        llm_tokenizer = AutoTokenizer.from_pretrained('facebook/galactica-1.3b', use_fast=False, padding_side='right')
+        plm_tokenizer = EsmTokenizer.from_pretrained('facebook/esm2_t30_150M_UR50D')
+        llm_tokenizer.add_special_tokens({'pad_token': '<pad>'})
+        loader = DataLoader(
+            train_dataset,
+            batch_size=32,
+            shuffle=True,
+            num_workers=4,
+            pin_memory=False,
+            drop_last=True,
+            persistent_workers=False,
+            collate_fn=ProtQACollater(llm_tokenizer, plm_tokenizer, 40, 40, 1024),
+        )
+    else:
+        print(len(train_dataset.data_list))
+        print(len(val_dataset.data_list))
+        print(len(test_dataset.data_list))
+        

data_provider/proteinchat_dm.py

@@ -0,0 +1,254 @@
+# Copyright (c) Microsoft Corporation.
+# Licensed under the MIT License.
+import json
+import random
+import torch
+import os
+import numpy as np
+from pytorch_lightning import LightningDataModule
+from torch.utils.data import DataLoader, Dataset
+from data_provider.gal_helpers import escape_custom_split_sequence
+from pathlib import Path
+from torch.utils.data.dataloader import default_collate
+
+
+class ProteinChatCollater(object):
+    def __init__(self, tokenizer, q_max_len, a_max_len, use_gal):
+        self.tokenizer = tokenizer
+        self.q_max_len = q_max_len
+        self.a_max_len = a_max_len
+        self.use_gal = use_gal
+        
+    def __call__(self, batch):
+        embeds, prot_seqs, questions, answers, q_types = zip(*batch)
+        max_embed_len = 896
+        ## concate 
+        if False:
+            max_dim = max([e.shape[0] for e in embeds])
+
+            padded_embeds = []
+            for embed in embeds:
+                shape_dim0 = embed.shape[0]
+                pad1 = ((0, max_dim - shape_dim0), (0, 0), (0, 0))
+                padded_embeds.append(np.pad(embed, pad1, mode='constant'))
+            padded_embeds = default_collate(padded_embeds).squeeze(dim=2)[:,:1024,:]
+        else:
+            padded_embeds = torch.zeros(len(embeds), max_embed_len, 512)
+            for i in range(len(embeds)):
+                padded_embeds[i, :embeds[i].shape[0], :] = embeds[i][:max_embed_len, :]
+            padded_embeds = padded_embeds.detach()
+
+        assert len(prot_seqs) == len(questions) == len(answers)
+
+        if self.use_gal:
+            questions = [escape_custom_split_sequence(q) for q in questions]
+        answers = [a + '\n' for a in answers]
+        self.tokenizer.padding_side = 'left'
+        q_batch = self.tokenizer(questions,
+                                 truncation=True,
+                                 padding='max_length',
+                                 add_special_tokens=True,
+                                 max_length=self.q_max_len,
+                                 return_tensors='pt',
+                                 return_attention_mask=True, 
+                                 return_token_type_ids=False)
+        self.tokenizer.padding_side = 'right'
+        a_batch = self.tokenizer(answers,
+                                 truncation=True,
+                                 padding='max_length',
+                                 add_special_tokens=True,
+                                 max_length=self.a_max_len,
+                                 return_tensors='pt',
+                                 return_attention_mask=True, 
+                                 return_token_type_ids=False)
+        prot_mask = torch.ones(padded_embeds.shape[0], padded_embeds.shape[1], dtype=torch.bool)
+        return (padded_embeds, prot_mask), q_batch, a_batch
+
+
+class InferenceCollater(object):
+    def __init__(self, tokenizer, q_max_len, a_max_len, use_gal):
+        self.tokenizer = tokenizer
+        self.q_max_len = q_max_len
+        self.a_max_len = a_max_len
+        self.use_gal = use_gal
+        
+    def __call__(self, batch):
+        embeds, prot_seqs, questions, answers, q_types = zip(*batch)
+        max_embed_len = 896
+        ## concate 
+        if False:
+            max_dim = max([e.shape[0] for e in embeds])
+
+            padded_embeds = []
+            for embed in embeds:
+                shape_dim0 = embed.shape[0]
+                pad1 = ((0, max_dim - shape_dim0), (0, 0), (0, 0))
+                padded_embeds.append(np.pad(embed, pad1, mode='constant'))
+            padded_embeds = default_collate(padded_embeds).squeeze(dim=2)[:,:1024,:]
+        else:
+            padded_embeds = torch.zeros(len(embeds), max_embed_len, 512)
+            for i in range(len(embeds)):
+                padded_embeds[i, :embeds[i].shape[0], :] = embeds[i][:max_embed_len, :]
+            padded_embeds = padded_embeds.detach()
+
+        assert len(prot_seqs) == len(questions) == len(answers)
+
+        if self.use_gal:
+            questions = [escape_custom_split_sequence(q) for q in questions]
+        answers = [a + '\n' for a in answers]
+        self.tokenizer.padding_side = 'left'
+        q_batch = self.tokenizer(questions,
+                                 truncation=True,
+                                 padding='max_length',
+                                 add_special_tokens=True,
+                                 max_length=self.q_max_len,
+                                 return_tensors='pt',
+                                 return_attention_mask=True, 
+                                 return_token_type_ids=False)
+        prot_mask = torch.ones(padded_embeds.shape[0], padded_embeds.shape[1], dtype=torch.bool)
+        target_dict = {'answers': answers, "q_types": q_types}
+        return (padded_embeds, prot_mask), q_batch, target_dict
+
+
+class ProteinChatDM(LightningDataModule):
+    def __init__(
+        self,
+        root: str = 'data/',
+        args=None,
+    ):
+        super().__init__()
+        self.args = args
+        self.batch_size = args.batch_size
+        self.inference_batch_size = args.inference_batch_size
+        self.num_workers = args.num_workers
+        self.q_max_len = args.q_max_len
+        self.a_max_len = args.a_max_len
+        self.prompt = args.prompt
+        
+        self.train_dataset = ProteinChatDataset(root, 'train.txt', prompt="### Human: {}\n### Assistant: ", pt_file_path=args.pt_file_path)
+        self.val_dataset = ProteinChatDataset(root, 'val.txt', prompt="### Human: {}\n### Assistant: ", pt_file_path=args.pt_file_path)
+        self.test_dataset = ProteinChatDataset(root, 'test.txt', prompt="### Human: {}\n### Assistant: ", pt_file_path=args.pt_file_path)
+        
+        self.tokenizer = None
+        self.use_gal = args.llm_name.find('gal') >= 0
+    
+    def init_tokenizer(self, tokenizer):
+        self.tokenizer = tokenizer
+
+    def train_dataloader(self):
+        loader = DataLoader(
+            self.train_dataset,
+            batch_size=self.batch_size,
+            shuffle=True,
+            num_workers=self.num_workers,
+            pin_memory=False,
+            drop_last=True,
+            persistent_workers=False,
+            collate_fn=ProteinChatCollater(self.tokenizer, self.q_max_len, self.a_max_len, self.use_gal),
+        )
+        return loader
+
+    def val_dataloader(self):
+        val_loader = DataLoader(
+            self.val_dataset,
+            batch_size=self.batch_size,
+            shuffle=False,
+            num_workers=self.num_workers,
+            pin_memory=False,
+            drop_last=False,
+            persistent_workers=False,
+            collate_fn=ProteinChatCollater(self.tokenizer, self.q_max_len, self.a_max_len, self.use_gal),
+        )
+        test_loader = DataLoader(
+            self.test_dataset,
+            batch_size=self.inference_batch_size,
+            shuffle=False,
+            num_workers=self.num_workers,
+            pin_memory=False,
+            drop_last=False,
+            persistent_workers=False,
+            collate_fn=InferenceCollater(self.tokenizer, self.q_max_len, self.a_max_len, self.use_gal),
+        )
+        return [val_loader, test_loader]
+    
+
+    def add_model_specific_args(parent_parser):
+        parser = parent_parser.add_argument_group("Data module")
+        parser.add_argument('--num_workers', type=int, default=2)
+        parser.add_argument('--batch_size', type=int, default=32)
+        parser.add_argument('--inference_batch_size', type=int, default=4)
+        parser.add_argument('--root', type=str, default='data/SwissProtV3')
+        parser.add_argument('--q_max_len', type=int, default=30)
+        parser.add_argument('--a_max_len', type=int, default=36)
+        parser.add_argument('--prompt', type=str, default='[START_AMINO]{}[END_AMINO]. Question: {} Answer:')
+        parser.add_argument('--pt_file_path', type=str, default='/home/XXXX-2/proteinchatdata/proteinchat')
+        return parent_parser
+
+
+
+class ProteinChatDataset(Dataset):
+    def __init__(self, root_path, subset, pt_file_path, prompt="Question: {} Answer:"):
+        super(ProteinChatDataset, self).__init__()
+        self.data_path = Path(root_path) / subset
+        self.qa_path = Path(root_path) / 'qa_all.json'
+        self.q_type_path = Path(root_path) / 'q_types.txt'
+        self.prompt = prompt
+
+        ## load dataset
+        with open(self.qa_path, 'r') as f:
+            qa_data = json.load(f)
+        
+        with open(self.data_path, 'r') as f:
+            lines = f.readlines()
+            pdb2seq = [line.strip().split('\t') for line in lines]
+        
+        ## process dataset
+        pdb_set = set(i[0] for i in pdb2seq)
+        ## filter qa data
+        qa_data = {k: v for k, v in qa_data.items() if k in pdb_set}
+        assert len(qa_data) == len(pdb_set), print(len(qa_data), len(pdb_set))
+        
+        pt_file = Path(pt_file_path).glob('*.pt')
+        pt_file_ids = {f.name.split('.pt')[0] for f in pt_file}
+        self.pt_file_path = pt_file_path
+
+        ## load q types
+        with open(self.q_type_path, 'r') as f:
+            q_types = [line.strip().split('\t') for line in f.readlines()]
+        self.q_type_dict = {q: t for q, t in q_types}
+
+        ## generate qa data
+        self.data_list = []
+        for pdb_id, seq in pdb2seq:
+            if pdb_id not in pt_file_ids:
+                continue
+            qa_list = qa_data[pdb_id]
+            for qa in qa_list:
+                q = qa['Q']
+                a = str(qa['A'])
+                self.data_list.append((pdb_id, seq, q, a))
+
+    def shuffle(self):
+        random.shuffle(self.data_list)
+        return self
+
+    def __len__(self):
+        return len(self.data_list)
+
+    def __getitem__(self, index):
+        pdb_id, seq, q, a = self.data_list[index]
+        q_type = self.q_type_dict[q]
+        path = os.path.join(self.pt_file_path, pdb_id + '.pt')
+        embed = torch.load(path, map_location=torch.device('cpu'))
+        embed = embed.squeeze(dim=1)
+        embed = embed.detach()
+        q = self.prompt.format(q)
+        return embed, seq, q, a, q_type
+
+
+if __name__ == '__main__':
+    dataset = ProteinChatDataset('./data/PDBDataset', 'train.txt')
+    dataset.shuffle()
+    for i in range(1000):
+        print(dataset[i][0].shape)
+    

data_provider/stage1_dm.py

@@ -0,0 +1,539 @@
+# Copyright (c) Microsoft Corporation.
+# Licensed under the MIT License.
+from pytorch_lightning import LightningDataModule
+import json
+from torch.utils.data import DataLoader, Dataset, ConcatDataset
+import random
+from pathlib import Path
+
+
+def rand_seq_crop(seq, max_len):
+    if len(seq) <= max_len:
+        return seq
+    rand_pos = random.randint(0, len(seq)-1-max_len)
+    return seq[rand_pos:rand_pos+max_len]
+
+
+class Stage1Collater(object):
+    def __init__(self, tokenizer, prot_tokenizer, text_max_len, prot_max_len, prot_aug='None'):
+        self.tokenizer = tokenizer
+        self.prot_tokenizer = prot_tokenizer
+        self.text_max_len = text_max_len
+        self.prot_max_len = prot_max_len
+        self.prot_aug = prot_aug
+    
+    def __call__(self, batch):
+        prot_seqs, text_seqs, _ = zip(*batch)
+        if self.prot_aug == 'rand_crop':
+            prot_seqs = [rand_seq_crop(seq, self.prot_max_len-2) for seq in prot_seqs] # -2 for the two special tokens
+        
+        text_tokens = self.tokenizer(text_seqs,
+                                     truncation=True,
+                                     padding='max_length',
+                                     add_special_tokens=True,
+                                     max_length=self.text_max_len,
+                                     return_tensors='pt',
+                                     return_attention_mask=True, 
+                                     return_token_type_ids=False)
+        prot_tokens = self.prot_tokenizer(prot_seqs,
+                                          truncation=True,
+                                          padding='max_length',
+                                          max_length=self.prot_max_len,
+                                          return_tensors="pt",
+                                          return_attention_mask=True, 
+                                          return_token_type_ids=False)
+        return prot_tokens, text_tokens
+
+
+class Stage1DM(LightningDataModule):
+    def __init__(
+        self,
+        num_workers: int = 0,
+        batch_size: int = 256,
+        root: str = 'data/',
+        args=None,
+    ):
+        super().__init__()
+        self.batch_size = batch_size
+        self.match_batch_size = args.match_batch_size
+        self.num_workers = num_workers
+        self.text_max_len = args.text_max_len
+        self.prot_max_len = args.prot_max_len
+        if root.find('SwissProt') >= 0:
+            self.train_dataset = SwissProtDataset(root+'/train_set.jsonl')
+            self.val_dataset = SwissProtDataset(root+'/valid_set.jsonl')
+            self.test_dataset = SwissProtDataset(root+'/test_set.jsonl')
+            self.val_dataset_match = SwissProtDataset(root+'/valid_set.jsonl').shuffle()
+            self.test_dataset_match = SwissProtDataset(root+'/test_set.jsonl').shuffle()    
+        elif root.find('PDBDataset') >= 0:
+            self.train_dataset = PDBAbstractDataset(root, 'train.txt')
+            self.val_dataset = PDBAbstractDataset(root, 'val.txt')
+            self.test_dataset = PDBAbstractDataset(root, 'test.txt')
+            self.val_dataset_match = PDBAbstractDataset(root, 'val.txt').shuffle()
+            self.test_dataset_match = PDBAbstractDataset(root, 'test.txt').shuffle()    
+        elif root.find('OntoProtein') >= 0:
+            self.train_dataset = OntoProteinDataset(root+'/train.txt')
+            self.val_dataset = OntoProteinDataset(root+'/valid.txt')
+            self.test_dataset = OntoProteinDataset(root+'/test.txt')
+            self.val_dataset_match = OntoProteinDataset(root+'/valid.txt').shuffle()
+            self.test_dataset_match = OntoProteinDataset(root+'/test.txt').shuffle()
+        else:
+            raise NotImplementedError
+
+        self.tokenizer = None
+        self.prot_tokenizer = None
+        self.prot_aug = args.prot_aug
+
+    def init_tokenizer(self, tokenizer, prot_tokenizer):
+        self.tokenizer = tokenizer
+        self.prot_tokenizer = prot_tokenizer
+
+    def train_dataloader(self):
+        loader = DataLoader(
+            self.train_dataset,
+            batch_size=self.batch_size,
+            shuffle=True,
+            num_workers=self.num_workers,
+            pin_memory=False,
+            drop_last=True,
+            # persistent_workers=True,
+            collate_fn=Stage1Collater(self.tokenizer, self.prot_tokenizer, self.text_max_len, self.prot_max_len, self.prot_aug)
+        )
+        return loader
+
+    def val_dataloader(self):
+        loader = DataLoader(
+            self.val_dataset,
+            batch_size=self.batch_size,
+            shuffle=False,
+            num_workers=self.num_workers,
+            pin_memory=False,
+            drop_last=False,
+            # persistent_workers=True,
+            collate_fn=Stage1Collater(self.tokenizer, self.prot_tokenizer, self.text_max_len, self.prot_max_len, self.prot_aug)
+        )
+        return loader
+
+    def match_dataloader(self):
+        val_match_loader = DataLoader(self.val_dataset_match, 
+                                      batch_size=self.match_batch_size,
+                                      shuffle=False,
+                                      num_workers=self.num_workers, 
+                                      pin_memory=False, 
+                                      drop_last=False, 
+                                    #   persistent_workers=True,
+                                      collate_fn=Stage1Collater(self.tokenizer, self.prot_tokenizer, self.text_max_len, self.prot_max_len, self.prot_aug))
+        test_match_loader = DataLoader(self.test_dataset_match, 
+                                       batch_size=self.match_batch_size,
+                                       shuffle=False,
+                                       num_workers=self.num_workers, 
+                                       pin_memory=False, 
+                                       drop_last=False, 
+                                    #    persistent_workers=True,
+                                       collate_fn=Stage1Collater(self.tokenizer, self.prot_tokenizer, self.text_max_len, self.prot_max_len, self.prot_aug))
+        return val_match_loader, test_match_loader
+
+    def add_model_specific_args(parent_parser):
+        parser = parent_parser.add_argument_group("Data module")
+        parser.add_argument('--num_workers', type=int, default=4)
+        parser.add_argument('--batch_size', type=int, default=32)
+        parser.add_argument('--match_batch_size', type=int, default=64)
+        parser.add_argument('--root', type=str, default='data/SwissProtV3')
+        parser.add_argument('--text_max_len', type=int, default=128)
+        parser.add_argument('--prot_max_len', type=int, default=1024)
+        parser.add_argument('--prot_aug', type=str, default='None')
+        return parent_parser
+    
+
+
+class Stage1MixDM(LightningDataModule):
+    def __init__(
+        self,
+        num_workers: int = 0,
+        batch_size: int = 256,
+        root: str = 'data/',
+        args=None,
+    ):
+        super().__init__()
+        self.batch_size = batch_size
+        self.match_batch_size = args.match_batch_size
+        self.num_workers = num_workers
+        self.text_max_len = args.text_max_len
+        self.prot_max_len = args.prot_max_len
+        assert args.mix_dataset
+      
+        train_dataset1 = SwissProtDataset(root+'/SwissProtV3/train_set.jsonl')
+        train_dataset2 = OntoProteinDataset(root+'/OntoProteinDatasetV2/train.txt')
+        # 新增 PDBAbstract 训练集（需指定 subset 为 'train.txt'）
+        train_dataset3 = PDBAbstractDataset(root + '/PDBDataset/', subset='train.txt')  
+
+        #self.train_dataset = ConcatDataset([train_dataset1, train_dataset2,train_dataset3], )
+        self.train_dataset = ConcatDataset([train_dataset1,train_dataset2], )
+
+        self.swiss_val_dataset = SwissProtDataset(root+'/SwissProtV3/valid_set.jsonl')
+        self.onto_val_dataset = OntoProteinDataset(root+'/OntoProteinDatasetV2/valid.txt')
+        self.pdb_val_dataset = PDBAbstractDataset(root + '/PDBDataset/', subset='val.txt')
+
+        self.swiss_test_dataset = SwissProtDataset(root+'/SwissProtV3/test_set.jsonl')
+        self.onto_test_dataset = OntoProteinDataset(root+'/OntoProteinDatasetV2/test.txt')
+        self.pdb_test_dataset = PDBAbstractDataset(root + '/PDBDataset/', subset='test.txt')
+        
+        self.swiss_val_dataset_match = SwissProtDataset(root+'/SwissProtV3/valid_set.jsonl').shuffle()
+        self.onto_val_dataset_match = OntoProteinDataset(root+'/OntoProteinDatasetV2/valid.txt').shuffle()
+        self.pdb_val_dataset_match = PDBAbstractDataset(root + '/PDBDataset/', subset='val.txt').shuffle()
+
+        self.swiss_test_dataset_match = SwissProtDataset(root+'/SwissProtV3/test_set.jsonl').shuffle()
+        self.onto_test_dataset_match = OntoProteinDataset(root+'/OntoProteinDatasetV2/test.txt').shuffle()
+        self.pdb_test_dataset_match = PDBAbstractDataset(root + '/PDBDataset/', subset='test.txt').shuffle()
+
+        self.tokenizer = None
+        self.prot_tokenizer = None
+        self.prot_aug = args.prot_aug
+
+    def init_tokenizer(self, tokenizer, prot_tokenizer):
+        self.tokenizer = tokenizer
+        self.prot_tokenizer = prot_tokenizer
+
+    def train_dataloader(self):
+        loader = DataLoader(
+            self.train_dataset,
+            batch_size=self.batch_size,
+            shuffle=True,
+            num_workers=self.num_workers,
+            pin_memory=False,
+            drop_last=True,
+            collate_fn=Stage1Collater(self.tokenizer, self.prot_tokenizer, self.text_max_len, self.prot_max_len, self.prot_aug)
+        )
+        return loader
+
+    def val_dataloader(self):
+        loader1 = DataLoader(
+            self.swiss_val_dataset,
+            batch_size=self.batch_size,
+            shuffle=False,
+            num_workers=self.num_workers,
+            pin_memory=False,
+            drop_last=False,
+            collate_fn=Stage1Collater(self.tokenizer, self.prot_tokenizer, self.text_max_len, self.prot_max_len, self.prot_aug)
+        )
+        loader2 = DataLoader(
+            self.onto_val_dataset,
+            batch_size=self.batch_size,
+            shuffle=False,
+            num_workers=self.num_workers,
+            pin_memory=False,
+            drop_last=False,
+            collate_fn=Stage1Collater(self.tokenizer, self.prot_tokenizer, self.text_max_len, self.prot_max_len, self.prot_aug)
+        )
+         # 新增 PDB 验证加载器
+        loader3 = DataLoader(
+            self.pdb_val_dataset,
+            batch_size=self.batch_size,
+            shuffle=False,
+            num_workers=self.num_workers,
+            pin_memory=False,
+            drop_last=False,
+            collate_fn=Stage1Collater(self.tokenizer, self.prot_tokenizer, self.text_max_len, self.prot_max_len, self.prot_aug)
+        )
+
+        return [loader1, loader2, loader3]
+
+    def swiss_match_dataloader(self):
+        val_match_loader = DataLoader(self.swiss_val_dataset_match,
+                                      batch_size=self.match_batch_size,
+                                      shuffle=False,
+                                      num_workers=self.num_workers, 
+                                      pin_memory=False, 
+                                      drop_last=False, 
+                                      collate_fn=Stage1Collater(self.tokenizer, self.prot_tokenizer, self.text_max_len, self.prot_max_len, self.prot_aug))
+        test_match_loader = DataLoader(self.swiss_test_dataset_match, 
+                                       batch_size=self.match_batch_size,
+                                       shuffle=False,
+                                       num_workers=self.num_workers, 
+                                       pin_memory=False, 
+                                       drop_last=False, 
+                                       collate_fn=Stage1Collater(self.tokenizer, self.prot_tokenizer, self.text_max_len, self.prot_max_len, self.prot_aug))
+        return val_match_loader, test_match_loader
+    
+    def onto_match_dataloader(self):
+        val_match_loader = DataLoader(self.onto_val_dataset_match,
+                                      batch_size=self.match_batch_size,
+                                      shuffle=False,
+                                      num_workers=self.num_workers, 
+                                      pin_memory=False, 
+                                      drop_last=False, 
+                                      collate_fn=Stage1Collater(self.tokenizer, self.prot_tokenizer, self.text_max_len, self.prot_max_len, self.prot_aug))
+        test_match_loader = DataLoader(self.onto_test_dataset_match, 
+                                       batch_size=self.match_batch_size,
+                                       shuffle=False,
+                                       num_workers=self.num_workers, 
+                                       pin_memory=False, 
+                                       drop_last=False, 
+                                       collate_fn=Stage1Collater(self.tokenizer, self.prot_tokenizer, self.text_max_len, self.prot_max_len, self.prot_aug))
+        return val_match_loader, test_match_loader
+    
+    def pdb_match_dataloader(self):
+        val_match_loader = DataLoader(
+            self.pdb_val_dataset_match,
+            batch_size=self.match_batch_size,
+            shuffle=False,
+            num_workers=self.num_workers,
+            pin_memory=False,
+            drop_last=False,
+            collate_fn=Stage1Collater(self.tokenizer, self.prot_tokenizer, self.text_max_len, self.prot_max_len, self.prot_aug)
+        )
+        test_match_loader = DataLoader(
+            self.pdb_test_dataset_match,
+            batch_size=self.match_batch_size,
+            shuffle=False,
+            num_workers=self.num_workers,
+            pin_memory=False,
+            drop_last=False,
+            collate_fn=Stage1Collater(self.tokenizer, self.prot_tokenizer, self.text_max_len, self.prot_max_len, self.prot_aug)
+        )
+        return val_match_loader, test_match_loader
+    
+    def add_model_specific_args(parent_parser):
+        parser = parent_parser.add_argument_group("Data module")
+        parser.add_argument('--num_workers', type=int, default=4)
+        parser.add_argument('--batch_size', type=int, default=32)
+        parser.add_argument('--match_batch_size', type=int, default=64)
+        parser.add_argument('--root', type=str, default='data')
+        parser.add_argument('--text_max_len', type=int, default=128)
+        parser.add_argument('--prot_max_len', type=int, default=1024)
+        parser.add_argument('--prot_aug', type=str, default='None')
+        return parent_parser
+
+
+# class SwissProtDataset(Dataset):
+#     def __init__(self, data_path, prompt='Swiss-Prot description: ', return_prompt=False):
+#         super(SwissProtDataset, self).__init__()
+#         self.data_path = data_path
+
+#         ## load data
+#         with open(data_path, 'r') as f:
+#             lines = f.readlines()
+#             lines = [line.strip() for line in lines]
+#             self.data_list = [json.loads(line) for line in lines]
+
+#         ## preprocessing
+#         self.data_list = [(p, t.strip() + '\n') for p, t in self.data_list]
+
+#         self.text2id = {}
+#         for prot_seq, text_seq in self.data_list:
+#             if text_seq not in self.text2id:
+#                 self.text2id[text_seq] = len(self.text2id)
+        
+#         self.prompt = prompt
+#         self.return_prompt = return_prompt
+
+#     def shuffle(self):
+#         random.shuffle(self.data_list)
+#         return self
+
+#     def len(self,):
+#         return len(self)
+    
+#     def get(self, idx):
+#         return self.__getitem__(idx)
+
+#     def __len__(self):
+#         return len(self.data_list)
+
+#     def __getitem__(self, index):
+#         prot_seq, text_seq = self.data_list[index]
+#         if self.return_prompt:
+#             return prot_seq, self.prompt, text_seq, index
+#         return prot_seq, text_seq, index
+
+class SwissProtDataset(Dataset):
+    def __init__(self, data_path, prompt='Swiss-Prot description: ', return_prompt=False):
+        super(SwissProtDataset, self).__init__()
+        self.data_path = data_path
+        self.prompt = prompt
+        self.return_prompt = return_prompt
+        
+        # 加载并预处理数据
+        self.data_list = self._load_and_preprocess(data_path)
+        self.text2id = self._build_text_vocab()
+
+    def _load_and_preprocess(self, data_path):
+        """加载JSONL文件并预处理"""
+        data_list = []
+        with open(data_path, 'r', encoding='utf-8') as f:
+            for line in f:
+                try:
+                    item = json.loads(line.strip())
+                    # 确保包含所需字段
+                    if 'protein' in item and 'text' in item:
+                        prot_seq = item['protein']
+                        text_seq = item['text'].strip() + '\n'  # 添加结尾换行符
+                        data_list.append((prot_seq, text_seq))
+                    else:
+                        print(f"警告: 跳过缺少字段的行: {line[:50]}...")
+                except json.JSONDecodeError:
+                    print(f"警告: 跳过无效JSON行: {line[:50]}...")
+        return data_list
+
+    def _build_text_vocab(self):
+        """构建文本到ID的映射"""
+        text2id = {}
+        for _, text_seq in self.data_list:
+            if text_seq not in text2id:
+                text2id[text_seq] = len(text2id)
+        return text2id
+
+    def shuffle(self):
+        """打乱数据集顺序"""
+        random.shuffle(self.data_list)
+        return self
+
+    def __len__(self):
+        return len(self.data_list)
+
+    def __getitem__(self, index):
+        prot_seq, text_seq = self.data_list[index]
+        if self.return_prompt:
+            return prot_seq, self.prompt, text_seq, index
+        return prot_seq, text_seq, index
+
+    # 添加一些实用方法
+    def get_protein_sequence(self, index):
+        """获取指定索引的蛋白质序列"""
+        return self.data_list[index][0]
+
+    def get_text_description(self, index):
+        """获取指定索引的文本描述"""
+        return self.data_list[index][1]
+
+    def get_text_id(self, text_seq):
+        """获取文本描述的ID"""
+        return self.text2id.get(text_seq, -1)
+
+
+class PDBAbstractDataset(Dataset):
+    def __init__(self, root_path, subset, prompt='ABSTRACT: ', return_prompt=False):
+        super(PDBAbstractDataset, self).__init__()
+        self.data_path = Path(root_path) / subset
+        self.abstract_path = Path(root_path) / 'abstract.json'
+        
+        ## load dataset
+        with open(self.abstract_path, 'r') as f:
+            abstract_data = json.load(f)
+            abstract_data_dict = {line['pdb_id']: line['caption'] for line in abstract_data}
+        
+        with open(self.data_path, 'r') as f:
+            lines = f.readlines()
+            pdb2seq = [line.strip().split('\t') for line in lines]
+        
+        ## process dataset
+        data_list = []
+        for pdb_id, seq in pdb2seq:
+            abstract = abstract_data_dict[pdb_id]
+            abstract = abstract.replace('\n', ' ').strip() + '\n'
+            data_list.append((seq, abstract))
+        self.data_list = data_list
+        self.prompt = prompt
+        self.return_prompt = return_prompt
+
+    def shuffle(self):
+        random.shuffle(self.data_list)
+        return self
+        
+    def len(self,):
+        return len(self)
+    
+    def get(self, idx):
+        return self.__getitem__(idx)
+
+    def __len__(self):
+        return len(self.data_list)
+
+    def __getitem__(self, index):
+        seq, abstract = self.data_list[index]
+        if self.return_prompt:
+            return seq, self.prompt, abstract, index
+        return seq, abstract,index
+    
+
+class OntoProteinDataset(Dataset):
+    def __init__(self, data_path, prompt='Gene Ontology description: ', return_prompt=False):
+        super(OntoProteinDataset, self).__init__()
+        self.data_path = data_path
+        
+        ## load data
+        with open(data_path, 'r') as f:
+            lines = f.readlines()
+            self.data_list = [line.strip().split('\t') for line in lines]
+
+        ## preprocessing
+        ## fixme: I have disabled the signal word for this dataset. However, it was used in previous experiments.
+        if True:
+            self.data_list = [(p, t.strip() + '\n') for p, t in self.data_list]
+        else:
+            self.data_list = [(p, "KG: " + t.strip() + '\n') for p, t in self.data_list]
+        self.prompt = prompt
+        self.return_prompt = return_prompt
+
+    def shuffle(self):
+        random.shuffle(self.data_list)
+        return self
+
+    def __len__(self):
+        return len(self.data_list)
+
+    def __getitem__(self, index):
+        prot_seq, text_seq = self.data_list[index]
+        if self.return_prompt:
+            return prot_seq, self.prompt, text_seq, index
+        return prot_seq, text_seq,index
+
+
+if __name__ == '__main__':
+    import numpy as np
+    ## get statistics for swiss prot dataset
+    if False:
+        swiss_train = SwissProtDataset('../data/SwissProtV3/train_set.jsonl')
+        swiss_valid = SwissProtDataset('../data/SwissProtV3/valid_set.jsonl')
+        swiss_test = SwissProtDataset('../data/SwissProtV3/test_set.jsonl')
+        print(len(swiss_train), len(swiss_valid), len(swiss_test))
+        
+        ## get amino acid statistics
+        aa_lens = np.asarray([len(seq) for seq, _ in swiss_train.data_list])
+        print('Train dataset mean: ', np.mean(aa_lens), 'min: ', aa_lens.min(), 'max: ', aa_lens.max())
+        aa_lens = np.asarray([len(seq) for seq, _ in swiss_valid.data_list])
+        print('Valid dataset mean: ', np.mean(aa_lens), 'min: ', aa_lens.min(), 'max: ', aa_lens.max())
+        aa_lens = np.asarray([len(seq) for seq, _ in swiss_test.data_list])
+        print('Test dataset mean: ', np.mean(aa_lens), 'min: ', aa_lens.min(), 'max: ', aa_lens.max())
+
+        ## get text statistics
+        text_lens = np.asarray([len(seq.split()) for _, seq in swiss_train.data_list])
+        print('Train dataset mean: ', np.mean(text_lens), 'min: ', text_lens.min(), 'max: ', text_lens.max())
+        text_lens = np.asarray([len(seq.split()) for _, seq in swiss_valid.data_list])
+        print('Valid dataset mean: ', np.mean(text_lens), 'min: ', text_lens.min(), 'max: ', text_lens.max())
+        text_lens = np.asarray([len(seq.split()) for _, seq in swiss_test.data_list])
+        print('Test dataset mean: ', np.mean(text_lens), 'min: ', text_lens.min(), 'max: ', text_lens.max())
+        print('---------------------------')
+
+    ## get statistics for onto protein dataset
+    onto_train = OntoProteinDataset('../data/OntoProteinDatasetV2/train.txt')
+    onto_valid = OntoProteinDataset('../data/OntoProteinDatasetV2/valid.txt')
+    onto_test = OntoProteinDataset('../data/OntoProteinDatasetV2/test.txt')
+    print(len(onto_train), len(onto_valid), len(onto_test))
+
+    ## get amino acid statistics
+    aa_lens = np.asarray([len(seq) for seq, _ in onto_train.data_list])
+    print('Train dataset mean: ', np.mean(aa_lens), 'min: ', aa_lens.min(), 'max: ', aa_lens.max())
+    aa_lens = np.asarray([len(seq) for seq, _ in onto_valid.data_list])
+    print('Valid dataset mean: ', np.mean(aa_lens), 'min: ', aa_lens.min(), 'max: ', aa_lens.max())
+    aa_lens = np.asarray([len(seq) for seq, _ in onto_test.data_list])
+    print('Test dataset mean: ', np.mean(aa_lens), 'min: ', aa_lens.min(), 'max: ', aa_lens.max())
+
+    ## get text statistics
+    text_lens = np.asarray([len(seq.split()) for _, seq in onto_train.data_list])
+    print('Train dataset mean: ', np.mean(text_lens), 'min: ', text_lens.min(), 'max: ', text_lens.max())
+    text_lens = np.asarray([len(seq.split()) for _, seq in onto_valid.data_list])
+    print('Valid dataset mean: ', np.mean(text_lens), 'min: ', text_lens.min(), 'max: ', text_lens.max())
+    text_lens = np.asarray([len(seq.split()) for _, seq in onto_test.data_list])
+    print('Test dataset mean: ', np.mean(text_lens), 'min: ', text_lens.min(), 'max: ', text_lens.max())
+    print('---------------------------')

data_provider/stage2_dm.py

@@ -0,0 +1,386 @@
+# Copyright (c) Microsoft Corporation.
+# Licensed under the MIT License.
+import json
+from pytorch_lightning import LightningDataModule
+from torch.utils.data import DataLoader, ConcatDataset
+from data_provider.stage1_dm import SwissProtDataset, OntoProteinDataset
+from data_provider.stage3_dm import DeepLocBinaryDataset,AlpacaDataset,MolInstructionDataset,DeepLocMultiDataset,Deepsol,DeepsoluE,Protsolm,FLIP_GB1,FLIP_AAV
+from data_provider.bindingdb import BindingDB
+from data_provider.metalIonbinding import MetallonBinding
+from data_provider.go import GO_BP,EC
+from data_provider.production import Antibiotic_Resistance,Thermostability,Material,Clone
+from data_provider.mutation import TAPE_Stability,TAPE_Fluorescence
+
+class Stage2Collater(object):
+    def __init__(self, tokenizer, prot_tokenizer, text_max_len, prot_max_len):
+        self.tokenizer = tokenizer
+        self.prot_tokenizer = prot_tokenizer
+        self.text_max_len = text_max_len
+        self.prot_max_len = prot_max_len
+        
+    def __call__(self, batch):
+        prot_seqs, prompt_seqs, text_seqs, _ = zip(*batch)
+        prot_tokens = self.prot_tokenizer(prot_seqs,
+                                          truncation=True,
+                                          padding='max_length',
+                                          max_length=self.prot_max_len,
+                                          return_tensors="pt",
+                                          return_attention_mask=True, 
+                                          return_token_type_ids=False)
+        if False:
+            self.tokenizer.padding_side = 'left'
+            prompt_tokens = self.tokenizer(prompt_seqs,
+                                        truncation=True,
+                                        padding='longest',
+                                        add_special_tokens=True,
+                                        max_length=self.text_max_len,
+                                        return_tensors='pt',
+                                        return_attention_mask=True, 
+                                        return_token_type_ids=False)
+            self.tokenizer.padding_side = 'right'
+            text_tokens = self.tokenizer(text_seqs,
+                                        truncation=True,
+                                        padding='max_length',
+                                        add_special_tokens=True,
+                                        max_length=self.text_max_len,
+                                        return_tensors='pt',
+                                        return_attention_mask=True, 
+                                        return_token_type_ids=False)
+        else:
+            self.tokenizer.padding_side = 'left'
+            prompt_tokens = self.tokenizer(prompt_seqs,
+                                           truncation=True,
+                                           padding='longest',
+                                           add_special_tokens=True,
+                                           max_length=self.text_max_len,
+                                           return_tensors='pt',
+                                           return_attention_mask=True, 
+                                           return_token_type_ids=False)
+            max_prompt_len = int(prompt_tokens.attention_mask.sum(dim=1).max())
+            input_pair = [[p, t] for p, t in zip(prompt_seqs, text_seqs)]
+            input_tokens = self.tokenizer(input_pair,
+                                          truncation=True,
+                                          padding='max_length',
+                                          add_special_tokens=True,
+                                          max_length=self.text_max_len + max_prompt_len,
+                                          return_tensors='pt',
+                                          return_attention_mask=True,
+                                          return_token_type_ids=True)
+        return prot_tokens, input_tokens
+
+
+class InferenceCollater(object):
+    def __init__(self, tokenizer, prot_tokenizer, text_max_len, prot_max_len):
+        self.tokenizer = tokenizer
+        self.prot_tokenizer = prot_tokenizer
+        self.text_max_len = text_max_len
+        self.prot_max_len = prot_max_len
+        
+    def __call__(self, batch):
+        prot_seqs, prompt_seqs, text_seqs, indices = zip(*batch)
+        # print("=========")
+        # print(prot_seqs)
+
+        self.tokenizer.padding_side = 'right'
+        prompt_tokens = self.tokenizer(prompt_seqs,
+                                       truncation=True,
+                                       padding='longest',
+                                       add_special_tokens=False,
+                                       max_length=self.text_max_len,
+                                       return_tensors='pt',
+                                       return_attention_mask=True, 
+                                       return_token_type_ids=False)
+        prot_tokens = self.prot_tokenizer(prot_seqs,
+                                          truncation=True,
+                                          padding='max_length',
+                                          max_length=self.prot_max_len,
+                                          return_tensors="pt",
+                                          return_attention_mask=True, 
+                                          return_token_type_ids=False)
+        target_dict = {'targets': text_seqs, 'indices': indices}
+        return prot_tokens, prompt_tokens, target_dict
+
+
+class Stage2DM(LightningDataModule):
+    def __init__(
+        self,
+        root: str = 'data/',
+        args=None,
+    ):
+        super().__init__()
+        self.args = args
+        self.batch_size = args.batch_size
+        self.inference_batch_size = args.inference_batch_size
+        self.num_workers = args.num_workers
+        self.text_max_len = args.text_max_len
+        self.prot_max_len = args.prot_max_len
+        self.prompt = args.prompt
+        
+          # self.train_dataset = AlpacaDataset('/nas/shared/kilab/wangyujia/pretrain_data/instruct/alpaca-gpt4-train.jsonl', prompt=self.prompt, return_prompt=True)
+        # self.val_dataset = AlpacaDataset('/nas/shared/kilab/wangyujia/pretrain_data/instruct/alpaca-gpt4-valid.jsonl', prompt=self.prompt, return_prompt=True)
+        # self.test_dataset = AlpacaDataset('/nas/shared/kilab/wangyujia/pretrain_data/instruct/alpaca-gpt4-test.jsonl', prompt=self.prompt, return_prompt=True)
+
+        # self.train_dataset = MolInstructionDataset('/oss/wangyujia/pretrain-bench/mol-instruction/train.jsonl', prompt='', return_prompt=True)
+        # self.val_dataset = MolInstructionDataset('/oss/wangyujia/pretrain-bench/mol-instruction/train.jsonl', prompt='', return_prompt=True)
+        # self.test_dataset = MolInstructionDataset('/oss/wangyujia/pretrain-bench/mol-instruction/train.jsonl', prompt='', return_prompt=True)
+     
+        if self.args.dataset=='deeplocbinary':
+            self.train_dataset = DeepLocBinaryDataset('/oss/wangyujia/pretrain-bench/locate/deeplocbinary/train.csv', prompt=self.prompt, return_prompt=True)
+            self.val_dataset = DeepLocBinaryDataset('/oss/wangyujia/pretrain-bench/locate/deeplocbinary/valid.csv', prompt=self.prompt, return_prompt=True)
+            self.test_dataset = DeepLocBinaryDataset('/oss/wangyujia/pretrain-bench/locate/deeplocbinary/test.csv', prompt=self.prompt, return_prompt=True)
+        
+        elif self.args.dataset=='deeplocmulti':
+            self.train_dataset = DeepLocMultiDataset('/oss/wangyujia/pretrain-bench/locate/deeplocmulti/train.csv',prompt=self.prompt, return_prompt=True)
+            self.val_dataset = DeepLocMultiDataset('/oss/wangyujia/pretrain-bench/locate/deeplocmulti/valid.csv', prompt=self.prompt, return_prompt=True)
+            self.test_dataset = DeepLocMultiDataset('/oss/wangyujia/pretrain-bench/locate/deeplocmulti/test.csv', prompt=self.prompt, return_prompt=True)
+        
+        elif self.args.dataset=='deepsol':
+            self.train_dataset = Deepsol('/nas/shared/kilab/wangyujia/sft_data/deepsol/clean/train.csv',prompt=self.prompt, return_prompt=True)
+            self.val_dataset = Deepsol('/nas/shared/kilab/wangyujia/sft_data/deepsol/clean/valid.csv', prompt=self.prompt, return_prompt=True)
+            self.test_dataset = Deepsol('/nas/shared/kilab/wangyujia/sft_data/deepsol/clean/test.csv', prompt=self.prompt, return_prompt=True)
+        elif self.args.dataset=='deepsolue':
+            self.train_dataset = DeepsoluE('/nas/shared/kilab/wangyujia/sft_data/deepsoluE/clean/train.csv',prompt=self.prompt, return_prompt=True)
+            self.val_dataset = DeepsoluE('/nas/shared/kilab/wangyujia/sft_data/deepsoluE/clean/test.csv', prompt=self.prompt, return_prompt=True)
+            self.test_dataset = DeepsoluE('/nas/shared/kilab/wangyujia/sft_data/deepsoluE/clean/test.csv', prompt=self.prompt, return_prompt=True)
+        elif self.args.dataset=='protsolm':
+            self.train_dataset = Protsolm('/oss/wangyujia/pretrain-bench/solubility/protsolm/train.csv',prompt=self.prompt, return_prompt=True)
+            self.val_dataset = Protsolm('/oss/wangyujia/pretrain-bench/solubility/protsolm/valid.csv', prompt=self.prompt, return_prompt=True)
+            self.test_dataset = Protsolm('/oss/wangyujia/pretrain-bench/solubility/protsolm/test.csv', prompt=self.prompt, return_prompt=True)
+        elif self.args.dataset=='gb1':
+            self.train_dataset = FLIP_GB1('/nas/shared/kilab/wangyujia/sft_data/mutation/gb1/train.csv',prompt=self.prompt, return_prompt=True)
+            self.val_dataset = FLIP_GB1('/nas/shared/kilab/wangyujia/sft_data/mutation/gb1/valid.csv', prompt=self.prompt, return_prompt=True)
+            self.test_dataset = FLIP_GB1('/nas/shared/kilab/wangyujia/sft_data/mutation/gb1/test.csv', prompt=self.prompt, return_prompt=True)
+        elif self.args.dataset=='gb1_low':
+            self.train_dataset = FLIP_GB1('/nas/shared/kilab/wangyujia/sft_data/mutation/gb1/clean/train.csv',prompt=self.prompt, return_prompt=True)
+            self.val_dataset = FLIP_GB1('/nas/shared/kilab/wangyujia/sft_data/mutation/gb1/clean/valid.csv', prompt=self.prompt, return_prompt=True)
+            self.test_dataset = FLIP_GB1('/nas/shared/kilab/wangyujia/sft_data/mutation/gb1/clean/test.csv', prompt=self.prompt, return_prompt=True)
+       
+        elif self.args.dataset=='aav':
+            self.train_dataset = FLIP_AAV('/nas/shared/kilab/wangyujia/sft_data/mutation/aav/train.csv',prompt=self.prompt, return_prompt=True)
+            self.val_dataset = FLIP_AAV('/nas/shared/kilab/wangyujia/sft_data/mutation/aav/valid.csv', prompt=self.prompt, return_prompt=True)
+            self.test_dataset = FLIP_AAV('/nas/shared/kilab/wangyujia/sft_data/mutation/aav/test.csv', prompt=self.prompt, return_prompt=True)
+        
+        elif self.args.dataset=='bindingdb':
+            self.train_dataset = BindingDB('/nas/shared/kilab/wangyujia/sft_data/bindingdb/clean/train_small.csv',prompt=self.prompt, return_prompt=True)
+            self.val_dataset = BindingDB('/nas/shared/kilab/wangyujia/sft_data/bindingdb/clean/valid_small.csv', prompt=self.prompt, return_prompt=True)
+            self.test_dataset = BindingDB('/nas/shared/kilab/wangyujia/sft_data/bindingdb/clean/test_small.csv', prompt=self.prompt, return_prompt=True)
+        
+        elif self.args.dataset=='metallonbinding':
+            self.train_dataset = MetallonBinding('/nas/shared/kilab/wangyujia/sft_data/MetalIonBinding/train.csv',prompt=self.prompt, return_prompt=True)
+            self.val_dataset = MetallonBinding('/nas/shared/kilab/wangyujia/sft_data/MetalIonBinding/valid.csv', prompt=self.prompt, return_prompt=True)
+            self.test_dataset = MetallonBinding('/nas/shared/kilab/wangyujia/sft_data/MetalIonBinding/test.csv', prompt=self.prompt, return_prompt=True)
+        
+        elif self.args.dataset=='bp':
+            self.train_dataset = GO_BP('/nas/shared/kilab/wangyujia/sft_data/go/clean/BP_train.csv',prompt=self.prompt, return_prompt=True)
+            self.val_dataset = GO_BP('/nas/shared/kilab/wangyujia/sft_data/go/clean/BP_valid.csv', prompt=self.prompt, return_prompt=True)
+            self.test_dataset = GO_BP('/nas/shared/kilab/wangyujia/sft_data/go/clean/BP_test.csv', prompt=self.prompt, return_prompt=True)
+        elif self.args.dataset=='ec':
+            self.train_dataset = EC('/nas/shared/kilab/wangyujia/sft_data/EC/train.csv',prompt=self.prompt, return_prompt=True)
+            self.val_dataset = EC('/nas/shared/kilab/wangyujia/sft_data/EC/valid.csv', prompt=self.prompt, return_prompt=True)
+            self.test_dataset = EC('/nas/shared/kilab/wangyujia/sft_data/EC/test.csv', prompt=self.prompt, return_prompt=True)
+
+        elif self.args.dataset=='antibiotic':
+            self.train_dataset = Antibiotic_Resistance('/nas/shared/kilab/wangyujia/sft_data/production/Antibiotic_Resistance/train.csv',prompt=self.prompt, return_prompt=True)
+            self.val_dataset = Antibiotic_Resistance('/nas/shared/kilab/wangyujia/sft_data/production/Antibiotic_Resistance/test.csv', prompt=self.prompt, return_prompt=True)
+            self.test_dataset = Antibiotic_Resistance('/nas/shared/kilab/wangyujia/sft_data/production/Antibiotic_Resistance/test.csv', prompt=self.prompt, return_prompt=True)
+        elif self.args.dataset=='thermostability':
+            self.train_dataset = Thermostability('/nas/shared/kilab/wangyujia/sft_data/production/Thermostability/train.csv',prompt=self.prompt, return_prompt=True)
+            self.val_dataset = Thermostability('/nas/shared/kilab/wangyujia/sft_data/production/Thermostability/valid.csv', prompt=self.prompt, return_prompt=True)
+            self.test_dataset = Thermostability('/nas/shared/kilab/wangyujia/sft_data/production/Thermostability/test.csv', prompt=self.prompt, return_prompt=True)
+        elif self.args.dataset=='material':
+            self.train_dataset = Material('/oss/wangyujia/ProtT3/ProtT3/data/sft/dataset/material_production/train.csv',prompt=self.prompt, return_prompt=True)
+            self.val_dataset = Material('/oss/wangyujia/ProtT3/ProtT3/data/sft/dataset/material_production/val.csv', prompt=self.prompt, return_prompt=True)
+            self.test_dataset = Material('/oss/wangyujia/ProtT3/ProtT3/data/sft/dataset/material_production/test.csv', prompt=self.prompt, return_prompt=True)
+        #6
+        elif self.args.dataset=='clone':
+            self.train_dataset = Clone('/oss/wangyujia/ProtT3/ProtT3/data/sft/dataset/cloning_clf/train.csv',prompt=self.prompt, return_prompt=True)
+            self.val_dataset = Clone('/oss/wangyujia/ProtT3/ProtT3/data/sft/dataset/cloning_clf/val.csv', prompt=self.prompt, return_prompt=True)
+            self.test_dataset = Clone('/oss/wangyujia/ProtT3/ProtT3/data/sft/dataset/cloning_clf/test.csv', prompt=self.prompt, return_prompt=True)
+        
+
+
+        elif self.args.dataset=='stability':
+            self.train_dataset = TAPE_Stability('/oss/wangyujia/ProtT3/ProtT3/data/sft/dataset/TAPE_Stability/train.csv',prompt=self.prompt, return_prompt=True)
+            self.val_dataset = TAPE_Stability('/oss/wangyujia/ProtT3/ProtT3/data/sft/dataset/TAPE_Stability/val.csv', prompt=self.prompt, return_prompt=True)
+            self.test_dataset = TAPE_Stability('/oss/wangyujia/ProtT3/ProtT3/data/sft/dataset/TAPE_Stability/test.csv', prompt=self.prompt, return_prompt=True)
+        elif self.args.dataset=='fluorescence':
+            self.train_dataset = TAPE_Fluorescence('/oss/wangyujia/ProtT3/ProtT3/data/sft/dataset/TAPE_Fluorescence/fluorescence_prediction_train.csv',prompt=self.prompt, return_prompt=True)
+            self.val_dataset = TAPE_Fluorescence('/oss/wangyujia/ProtT3/ProtT3/data/sft/dataset/TAPE_Fluorescence/fluorescence_prediction_valid.csv', prompt=self.prompt, return_prompt=True)
+            self.test_dataset = TAPE_Fluorescence('/oss/wangyujia/ProtT3/ProtT3/data/sft/dataset/TAPE_Fluorescence/test.csv', prompt=self.prompt, return_prompt=True)
+        
+        
+        elif self.args.dataset=='empty':
+            self.train_dataset = Empty( prompt=self.prompt, return_prompt=True)
+            self.val_dataset = Empty( prompt=self.prompt, return_prompt=True)
+            self.test_dataset = Empty(prompt=self.prompt, return_prompt=True)
+        
+        elif self.args.dataset=='swiss-prot':
+            self.train_dataset = SwissProtDataset(root+'/SwissProtV3/train_set.jsonl', prompt='Swiss-Prot description: ', return_prompt=True)
+            self.val_dataset = SwissProtDataset(root+'/SwissProtV3/valid_set.jsonl', prompt='Swiss-Prot description: ', return_prompt=True)
+            self.test_dataset = SwissProtDataset(root+'/SwissProtV3/test_set.jsonl', prompt='Swiss-Prot description: ', return_prompt=True)
+
+        self.tokenizer = None
+        self.prot_tokenizer = None
+    
+    def init_tokenizer(self, tokenizer, prot_tokenizer):
+        self.tokenizer = tokenizer
+        self.prot_tokenizer = prot_tokenizer
+
+    def train_dataloader(self):
+        loader = DataLoader(
+            self.train_dataset,
+            batch_size=self.batch_size,
+            shuffle=True,
+            num_workers=self.num_workers,
+            pin_memory=False,
+            drop_last=True,
+            persistent_workers=False,
+            collate_fn=InferenceCollater(self.tokenizer, self.prot_tokenizer, self.text_max_len, self.prot_max_len),
+        )
+        return loader
+
+    def val_dataloader(self):
+        val_loader = DataLoader(
+            self.val_dataset,
+            batch_size=self.batch_size,
+            shuffle=False,
+            num_workers=self.num_workers,
+            pin_memory=False,
+            drop_last=False,
+            persistent_workers=False,
+            collate_fn=InferenceCollater(self.tokenizer, self.prot_tokenizer, self.text_max_len, self.prot_max_len),
+        )
+        test_loader = DataLoader(
+            self.test_dataset,
+            batch_size=self.inference_batch_size,
+            shuffle=False,
+            num_workers=self.num_workers,
+            pin_memory=False,
+            drop_last=False,
+            persistent_workers=False,
+            collate_fn=InferenceCollater(self.tokenizer, self.prot_tokenizer, self.text_max_len, self.prot_max_len),
+        )
+        return [val_loader, test_loader]
+    
+
+    def add_model_specific_args(parent_parser):
+        parser = parent_parser.add_argument_group("Data module")
+        parser.add_argument('--num_workers', type=int, default=2)
+        parser.add_argument('--batch_size', type=int, default=32)
+        parser.add_argument('--inference_batch_size', type=int, default=4)
+        parser.add_argument('--root', type=str, default='data')
+        parser.add_argument('--text_max_len', type=int, default=2048)
+        parser.add_argument('--q_max_len', type=int, default=29)
+        parser.add_argument('--a_max_len', type=int, default=36)
+        parser.add_argument('--prot_max_len', type=int, default=1024)
+        parser.add_argument('--prompt', type=str, default='The protein has the following properties:')
+        parser.add_argument('--filter_side_qa', action='store_true', default=False)
+        return parent_parser
+
+
+
+class Stage2MixDM(LightningDataModule):
+    def __init__(
+        self,
+        root: str = 'data/',
+        args=None,
+    ):
+        super().__init__()
+        self.args = args
+        self.batch_size = args.batch_size
+        self.inference_batch_size = args.inference_batch_size
+        self.num_workers = args.num_workers
+        self.text_max_len = args.text_max_len
+        self.prot_max_len = args.prot_max_len
+        # self.prompt = args.prompt
+        assert args.mix_dataset
+        
+        train_dataset1 = SwissProtDataset(root+'/SwissProtV3/train_set.jsonl', prompt='Swiss-Prot description: ', return_prompt=True)
+        train_dataset2 = OntoProteinDataset(root+'/OntoProteinDatasetV2/train.txt', prompt='Gene Ontology description: ', return_prompt=True)
+        self.train_dataset = ConcatDataset([train_dataset1,train_dataset2])
+        self.swiss_val_dataset = SwissProtDataset(root+'/SwissProtV3/valid_set.jsonl', prompt='Swiss-Prot description: ', return_prompt=True)
+        self.onto_val_dataset = OntoProteinDataset(root+'/OntoProteinDatasetV2/valid.txt', prompt='Gene Ontology description: ', return_prompt=True)
+        self.swiss_test_dataset = SwissProtDataset(root+'/SwissProtV3/test_set.jsonl', prompt='Swiss-Prot description: ', return_prompt=True)
+        self.onto_test_dataset = OntoProteinDataset(root+'/OntoProteinDatasetV2/test.txt', prompt='Gene Ontology description: ', return_prompt=True)
+        
+        
+
+        self.tokenizer = None
+        self.prot_tokenizer = None
+    
+    def init_tokenizer(self, tokenizer, prot_tokenizer):
+        self.tokenizer = tokenizer
+        self.prot_tokenizer = prot_tokenizer
+
+    def train_dataloader(self):
+        loader = DataLoader(
+            self.train_dataset,
+            batch_size=self.batch_size,
+            shuffle=True,
+            num_workers=self.num_workers,
+            pin_memory=False,
+            drop_last=True,
+            persistent_workers=False,
+            collate_fn=InferenceCollater(self.tokenizer, self.prot_tokenizer, self.text_max_len, self.prot_max_len),
+        )
+        return loader
+
+    def val_dataloader(self):
+        swiss_val_loader = DataLoader(
+            self.swiss_val_dataset,
+            batch_size=self.batch_size,
+            shuffle=False,
+            num_workers=self.num_workers,
+            pin_memory=False,
+            drop_last=False,
+            persistent_workers=False,
+            collate_fn=InferenceCollater(self.tokenizer, self.prot_tokenizer, self.text_max_len, self.prot_max_len),
+        )
+        swiss_test_loader = DataLoader(
+            self.swiss_test_dataset,
+            batch_size=self.inference_batch_size,
+            shuffle=False,
+            num_workers=self.num_workers,
+            pin_memory=False,
+            drop_last=False,
+            persistent_workers=False,
+            collate_fn=InferenceCollater(self.tokenizer, self.prot_tokenizer, self.text_max_len, self.prot_max_len),
+        )
+
+        onto_val_loader = DataLoader(
+            self.onto_val_dataset,
+            batch_size=self.batch_size,
+            shuffle=False,
+            num_workers=self.num_workers,
+            pin_memory=False,
+            drop_last=False,
+            persistent_workers=False,
+            collate_fn=InferenceCollater(self.tokenizer, self.prot_tokenizer, self.text_max_len, self.prot_max_len),
+        )
+        onto_test_loader = DataLoader(
+            self.onto_test_dataset,
+            batch_size=self.inference_batch_size,
+            shuffle=False,
+            num_workers=self.num_workers,
+            pin_memory=False,
+            drop_last=False,
+            persistent_workers=False,
+            collate_fn=InferenceCollater(self.tokenizer, self.prot_tokenizer, self.text_max_len, self.prot_max_len),
+        )
+        return [swiss_val_loader, swiss_test_loader, onto_val_loader, onto_test_loader]
+    
+
+    def add_model_specific_args(parent_parser):
+        parser = parent_parser.add_argument_group("Data module")
+        parser.add_argument('--num_workers', type=int, default=2)
+        parser.add_argument('--batch_size', type=int, default=32)
+        parser.add_argument('--inference_batch_size', type=int, default=4)
+        parser.add_argument('--root', type=str, default='data')
+        parser.add_argument('--text_max_len', type=int, default=1024)
+        parser.add_argument('--q_max_len', type=int, default=29)
+        parser.add_argument('--a_max_len', type=int, default=36)
+        parser.add_argument('--prot_max_len', type=int, default=1024)
+        # parser.add_argument('--prompt', type=str, default='The protein has the following properties: ')
+        return parent_parser
+
+