diff --git a/claude-1000/-home-ubuntu-ppi-project/5a87a0cd-6802-4ac5-a708-29b2bf7ea563/tasks/a07e7799acd8be342.output b/claude-1000/-home-ubuntu-ppi-project/5a87a0cd-6802-4ac5-a708-29b2bf7ea563/tasks/a07e7799acd8be342.output new file mode 100644 index 0000000000000000000000000000000000000000..4efa8fe8dc017bb7611776499e39897c7511c258 --- /dev/null +++ b/claude-1000/-home-ubuntu-ppi-project/5a87a0cd-6802-4ac5-a708-29b2bf7ea563/tasks/a07e7799acd8be342.output @@ -0,0 +1,55 @@ +{"parentUuid":null,"isSidechain":true,"promptId":"ae94540a-045a-4847-b16d-c2358b5c9db6","agentId":"a07e7799acd8be342","type":"user","message":{"role":"user","content":"In /home/ubuntu/ppi_project, I need to understand the data loading and training pipeline. Please find and summarize:\n1. Dataset/dataloader files (how samples are loaded, what fields each sample has — sequences, contact maps, residue indices, etc.)\n2. Where sequences get tokenized or padded\n3. Where contact maps are used (shape, format)\n4. The training loop or collate function that assembles batches\n5. Any existing cropping or windowing logic\n\nSearch thoroughly — look at files like dataset.py, data.py, train.py, collate, model.py, etc. Report file paths, key class/function names, and the shape/format of the data fields that flow through the pipeline."},"uuid":"ecb9c350-e516-4247-9e3e-a9dd81f0dda5","timestamp":"2026-06-01T15:28:16.877Z","userType":"external","entrypoint":"claude-vscode","cwd":"/home/ubuntu/ppi_project","sessionId":"5a87a0cd-6802-4ac5-a708-29b2bf7ea563","version":"2.1.159","gitBranch":"afdb"} +{"parentUuid":"ecb9c350-e516-4247-9e3e-a9dd81f0dda5","isSidechain":true,"agentId":"a07e7799acd8be342","attachment":{"type":"skill_listing","content":"- deep-research: Deep research harness — fan-out web searches, fetch sources, adversarially verify claims, synthesize a cited report. - When the user wants a deep, multi-source, fact-checked research report on any topic. BEFORE invoking, check if the question is specific enough to research directly — if underspecified (e.g., \"what car to buy\" without budget/use-case/region), ask 2-3 clarifying questions to narrow scope. Then pass the refined question as args, weaving the answers in.\n- update-config: Use this skill to configure the Claude Code harness via settings.json. Automated behaviors (\"from now on when X\", \"each time X\", \"whenever X\", \"before/after X\") require hooks configured in settings.json - the harness executes these, not Claude, so memory/preferences cannot fulfill them. Also use for: permissions (\"allow X\", \"add permission\", \"move permission to\"), env vars (\"set X=Y\"), hook troubleshooting, or any changes to settings.json/settings.local.json files. Examples: \"allow npm commands\", \"add bq permission to global settings\", \"move permission to user settings\", \"set DEBUG=true\", \"when claude stops show X\". For simple settings like theme/model, suggest the /config command.\n- keybindings-help: Use when the user wants to customize keyboard shortcuts, rebind keys, add chord bindings, or modify ~/.claude/keybindings.json. Examples: \"rebind ctrl+s\", \"add a chord shortcut\", \"change the submit key\", \"customize keybindings\".\n- verify: Verify that a code change actually does what it's supposed to by running the app and observing behavior. Use when asked to verify a PR, confirm a fix works, test a change manually, check that a feature works, or validate local changes before pushing.\n- code-review: Review the current diff for correctness bugs and reuse/simplification/efficiency cleanups at the given effort level (low/medium: fewer, high-confidence findings; high→max: broader coverage, may include uncertain findings). Pass --comment to post findings as inline PR comments, or --fix to apply the findings to the working tree after the review.\n- simplify: Review the changed code for reuse, simplification, efficiency, and altitude cleanups, then apply the fixes. Quality only — it does not hunt for bugs; use /code-review for that.\n- fewer-permission-prompts: Scan your transcripts for common read-only Bash and MCP tool calls, then add a prioritized allowlist to project .claude/settings.json to reduce permission prompts.\n- loop: Run a prompt or slash command on a recurring interval (e.g. /loop 5m /foo, defaults to 10m) - When the user wants to set up a recurring task, poll for status, or run something repeatedly on an interval (e.g. \"check the deploy every 5 minutes\", \"keep running /babysit-prs\"). Do NOT invoke for one-off tasks.\n- schedule: Create, update, list, or run scheduled remote agents (routines) that execute on a cron schedule. - When the user wants to schedule a recurring remote agent, set up automated tasks, create a cron job for Claude Code, or manage their scheduled agents/routines. Also use when the user wants a one-time scheduled run (\"run this once at 3pm\", \"remind me to check X tomorrow\").\n- claude-api: Build, debug, and optimize Claude API / Anthropic SDK apps. Apps built with this skill should include prompt caching. Also handles migrating existing Claude API code between Claude model versions (4.5 → 4.6, 4.6 → 4.7, retired-model replacements).\nTRIGGER when: code imports `anthropic`/`@anthropic-ai/sdk`; user asks for the Claude API, Anthropic SDK, or Managed Agents; user adds/modifies/tunes a Claude feature (caching, thinking, compaction, tool use, batch, files, citations, memory) or model (Opus/Sonnet/Haiku) in a file; questions about prompt caching / cache hit rate in an Anthropic SDK project.\nSKIP: file imports `openai`/other-provider SDK, filename like `*-openai.py`/`*-generic.py`, provider-neutral code, general programming/ML.\n- run: Launch and drive this project's app to see a change working. Use when asked to run, start, or screenshot the app, or to confirm a change works in the real app (not just tests). First looks for a project skill that already covers launching the app; otherwise falls back to built-in patterns per project type (CLI, server, TUI, Electron, browser-driven, library).\n- init: Initialize a new CLAUDE.md file with codebase documentation\n- review: Review a pull request\n- security-review: Complete a security review of the pending changes on the current branch","skillCount":14,"isInitial":true,"names":["deep-research","update-config","keybindings-help","verify","code-review","simplify","fewer-permission-prompts","loop","schedule","claude-api","run","init","review","security-review"]},"type":"attachment","uuid":"998ccbce-c076-41a1-b313-f70ddfe647a1","timestamp":"2026-06-01T15:28:16.883Z","userType":"external","entrypoint":"claude-vscode","cwd":"/home/ubuntu/ppi_project","sessionId":"5a87a0cd-6802-4ac5-a708-29b2bf7ea563","version":"2.1.159","gitBranch":"afdb"} +{"parentUuid":"998ccbce-c076-41a1-b313-f70ddfe647a1","isSidechain":true,"agentId":"a07e7799acd8be342","message":{"model":"claude-haiku-4-5-20251001","id":"msg_01Nz1AhiAWpZN4mGEqcKByby","type":"message","role":"assistant","content":[{"type":"text","text":"I'll conduct a thorough search of the codebase to understand the data loading and training pipeline. Let me start by exploring the directory structure and finding relevant files."}],"stop_reason":null,"stop_sequence":null,"stop_details":null,"usage":{"input_tokens":3,"cache_creation_input_tokens":25418,"cache_read_input_tokens":0,"cache_creation":{"ephemeral_5m_input_tokens":25418,"ephemeral_1h_input_tokens":0},"output_tokens":3,"service_tier":"standard","inference_geo":"not_available"},"diagnostics":null},"requestId":"req_011Cbcmz5MGz4FBdHaHxTiXG","attributionAgent":"Explore","type":"assistant","uuid":"2fddd289-0bf8-4775-a014-26fe5da38bc1","timestamp":"2026-06-01T15:28:18.961Z","userType":"external","entrypoint":"claude-vscode","cwd":"/home/ubuntu/ppi_project","sessionId":"5a87a0cd-6802-4ac5-a708-29b2bf7ea563","version":"2.1.159","gitBranch":"afdb"} +{"parentUuid":"2fddd289-0bf8-4775-a014-26fe5da38bc1","isSidechain":true,"agentId":"a07e7799acd8be342","message":{"model":"claude-haiku-4-5-20251001","id":"msg_01Nz1AhiAWpZN4mGEqcKByby","type":"message","role":"assistant","content":[{"type":"tool_use","id":"toolu_01PiW9dbu36SeTrQZJFnoQUE","name":"Bash","input":{"command":"find /home/ubuntu/ppi_project -type f -name \"*.py\" | grep -E \"(dataset|data|train|collate|model)\" | head -30"},"caller":{"type":"direct"}}],"stop_reason":null,"stop_sequence":null,"stop_details":null,"usage":{"input_tokens":3,"cache_creation_input_tokens":25418,"cache_read_input_tokens":0,"cache_creation":{"ephemeral_5m_input_tokens":25418,"ephemeral_1h_input_tokens":0},"output_tokens":3,"service_tier":"standard","inference_geo":"not_available"},"diagnostics":null},"requestId":"req_011Cbcmz5MGz4FBdHaHxTiXG","attributionAgent":"Explore","type":"assistant","uuid":"62c2e7dd-1e6b-43d1-a18e-4e89e722c819","timestamp":"2026-06-01T15:28:19.062Z","userType":"external","entrypoint":"claude-vscode","cwd":"/home/ubuntu/ppi_project","sessionId":"5a87a0cd-6802-4ac5-a708-29b2bf7ea563","version":"2.1.159","gitBranch":"afdb"} +{"parentUuid":"62c2e7dd-1e6b-43d1-a18e-4e89e722c819","isSidechain":true,"promptId":"ae94540a-045a-4847-b16d-c2358b5c9db6","agentId":"a07e7799acd8be342","type":"user","message":{"role":"user","content":[{"tool_use_id":"toolu_01PiW9dbu36SeTrQZJFnoQUE","type":"tool_result","content":"/home/ubuntu/ppi_project/evaluate_rf2_benchmark_2model.py\n/home/ubuntu/ppi_project/model.py\n/home/ubuntu/ppi_project/push_model_to_hf.py\n/home/ubuntu/ppi_project/train.py\n/home/ubuntu/ppi_project/dataset.py\n/home/ubuntu/ppi_project/scripts/cluster_dataset.py\n/home/ubuntu/ppi_project/hf_model/__init__.py\n/home/ubuntu/ppi_project/hf_model/configuration_flashppi.py\n/home/ubuntu/ppi_project/hf_model/modeling_flashppi.py\n/home/ubuntu/ppi_project/plm_models/modeling_esm_flash.py\n/home/ubuntu/ppi_project/plm_models/configuration_glm2.py\n/home/ubuntu/ppi_project/plm_models/modeling_glm2_flash.py\n/home/ubuntu/ppi_project/ppi-env/lib/python3.12/site-packages/charset_normalizer/models.py\n/home/ubuntu/ppi_project/ppi-env/lib/python3.12/site-packages/stack_data/utils.py\n/home/ubuntu/ppi_project/ppi-env/lib/python3.12/site-packages/stack_data/serializing.py\n/home/ubuntu/ppi_project/ppi-env/lib/python3.12/site-packages/stack_data/core.py\n/home/ubuntu/ppi_project/ppi-env/lib/python3.12/site-packages/stack_data/__init__.py\n/home/ubuntu/ppi_project/ppi-env/lib/python3.12/site-packages/stack_data/formatting.py\n/home/ubuntu/ppi_project/ppi-env/lib/python3.12/site-packages/stack_data/version.py\n/home/ubuntu/ppi_project/ppi-env/lib/python3.12/site-packages/pydantic/root_model.py\n/home/ubuntu/ppi_project/ppi-env/lib/python3.12/site-packages/pydantic/dataclasses.py\n/home/ubuntu/ppi_project/ppi-env/lib/python3.12/site-packages/pyarrow/dataset.py\n/home/ubuntu/ppi_project/ppi-env/lib/python3.12/site-packages/tzdata/__init__.py\n/home/ubuntu/ppi_project/ppi-env/lib/python3.12/site-packages/transformers/modeling_attn_mask_utils.py\n/home/ubuntu/ppi_project/ppi-env/lib/python3.12/site-packages/transformers/modeling_flash_attention_utils.py\n/home/ubuntu/ppi_project/ppi-env/lib/python3.12/site-packages/transformers/modeling_flax_outputs.py\n/home/ubuntu/ppi_project/ppi-env/lib/python3.12/site-packages/transformers/modeling_layers.py\n/home/ubuntu/ppi_project/ppi-env/lib/python3.12/site-packages/transformers/trainer_utils.py\n/home/ubuntu/ppi_project/ppi-env/lib/python3.12/site-packages/transformers/model_debugging_utils.py\n/home/ubuntu/ppi_project/ppi-env/lib/python3.12/site-packages/transformers/modeling_utils.py","is_error":false}]},"uuid":"20065174-fbb2-4a48-a06e-4cc757456cd6","timestamp":"2026-06-01T15:28:19.171Z","sourceToolAssistantUUID":"62c2e7dd-1e6b-43d1-a18e-4e89e722c819","userType":"external","entrypoint":"claude-vscode","cwd":"/home/ubuntu/ppi_project","sessionId":"5a87a0cd-6802-4ac5-a708-29b2bf7ea563","version":"2.1.159","gitBranch":"afdb"} +{"parentUuid":"20065174-fbb2-4a48-a06e-4cc757456cd6","isSidechain":true,"agentId":"a07e7799acd8be342","message":{"model":"claude-haiku-4-5-20251001","id":"msg_01Y3KAX9SZFEFyV6QPJRvkTf","type":"message","role":"assistant","content":[{"type":"text","text":"Good. Now let me read the main files related to data loading and training:"}],"stop_reason":null,"stop_sequence":null,"stop_details":null,"usage":{"input_tokens":5,"cache_creation_input_tokens":1125,"cache_read_input_tokens":25418,"cache_creation":{"ephemeral_5m_input_tokens":1125,"ephemeral_1h_input_tokens":0},"output_tokens":1,"service_tier":"standard","inference_geo":"not_available"},"diagnostics":null},"requestId":"req_011CbcmzF4vSVEL4rE83jzz9","attributionAgent":"Explore","type":"assistant","uuid":"651cc62c-9047-4158-99f6-f2293f632725","timestamp":"2026-06-01T15:28:20.589Z","userType":"external","entrypoint":"claude-vscode","cwd":"/home/ubuntu/ppi_project","sessionId":"5a87a0cd-6802-4ac5-a708-29b2bf7ea563","version":"2.1.159","gitBranch":"afdb"} +{"parentUuid":"651cc62c-9047-4158-99f6-f2293f632725","isSidechain":true,"agentId":"a07e7799acd8be342","message":{"model":"claude-haiku-4-5-20251001","id":"msg_01Y3KAX9SZFEFyV6QPJRvkTf","type":"message","role":"assistant","content":[{"type":"tool_use","id":"toolu_01FiKmUuqesxxTt8QqgLKUYa","name":"Read","input":{"file_path":"/home/ubuntu/ppi_project/dataset.py"},"caller":{"type":"direct"}}],"stop_reason":null,"stop_sequence":null,"stop_details":null,"usage":{"input_tokens":5,"cache_creation_input_tokens":1125,"cache_read_input_tokens":25418,"cache_creation":{"ephemeral_5m_input_tokens":1125,"ephemeral_1h_input_tokens":0},"output_tokens":1,"service_tier":"standard","inference_geo":"not_available"},"diagnostics":null},"requestId":"req_011CbcmzF4vSVEL4rE83jzz9","attributionAgent":"Explore","type":"assistant","uuid":"b3bf75ea-639e-4193-9cf7-f47b8a9939c0","timestamp":"2026-06-01T15:28:20.829Z","userType":"external","entrypoint":"claude-vscode","cwd":"/home/ubuntu/ppi_project","sessionId":"5a87a0cd-6802-4ac5-a708-29b2bf7ea563","version":"2.1.159","gitBranch":"afdb"} +{"parentUuid":"b3bf75ea-639e-4193-9cf7-f47b8a9939c0","isSidechain":true,"agentId":"a07e7799acd8be342","message":{"model":"claude-haiku-4-5-20251001","id":"msg_01Y3KAX9SZFEFyV6QPJRvkTf","type":"message","role":"assistant","content":[{"type":"tool_use","id":"toolu_01T9fSVBDYvoikLQDCLeoNwd","name":"Read","input":{"file_path":"/home/ubuntu/ppi_project/train.py"},"caller":{"type":"direct"}}],"stop_reason":null,"stop_sequence":null,"stop_details":null,"usage":{"input_tokens":5,"cache_creation_input_tokens":1125,"cache_read_input_tokens":25418,"cache_creation":{"ephemeral_5m_input_tokens":1125,"ephemeral_1h_input_tokens":0},"output_tokens":1,"service_tier":"standard","inference_geo":"not_available"},"diagnostics":null},"requestId":"req_011CbcmzF4vSVEL4rE83jzz9","attributionAgent":"Explore","type":"assistant","uuid":"a73d39f6-30fb-4a2d-9cfb-1f2bc30d10e4","timestamp":"2026-06-01T15:28:21.041Z","userType":"external","entrypoint":"claude-vscode","cwd":"/home/ubuntu/ppi_project","sessionId":"5a87a0cd-6802-4ac5-a708-29b2bf7ea563","version":"2.1.159","gitBranch":"afdb"} +{"parentUuid":"a73d39f6-30fb-4a2d-9cfb-1f2bc30d10e4","isSidechain":true,"agentId":"a07e7799acd8be342","message":{"model":"claude-haiku-4-5-20251001","id":"msg_01Y3KAX9SZFEFyV6QPJRvkTf","type":"message","role":"assistant","content":[{"type":"tool_use","id":"toolu_016aL1jZdYhhjwF8otCoAhR2","name":"Read","input":{"file_path":"/home/ubuntu/ppi_project/model.py"},"caller":{"type":"direct"}}],"stop_reason":null,"stop_sequence":null,"stop_details":null,"usage":{"input_tokens":5,"cache_creation_input_tokens":1125,"cache_read_input_tokens":25418,"cache_creation":{"ephemeral_5m_input_tokens":1125,"ephemeral_1h_input_tokens":0},"output_tokens":1,"service_tier":"standard","inference_geo":"not_available"},"diagnostics":null},"requestId":"req_011CbcmzF4vSVEL4rE83jzz9","attributionAgent":"Explore","type":"assistant","uuid":"b88a253a-19a3-42d0-932f-4e91525e7939","timestamp":"2026-06-01T15:28:21.119Z","userType":"external","entrypoint":"claude-vscode","cwd":"/home/ubuntu/ppi_project","sessionId":"5a87a0cd-6802-4ac5-a708-29b2bf7ea563","version":"2.1.159","gitBranch":"afdb"} +{"parentUuid":"b3bf75ea-639e-4193-9cf7-f47b8a9939c0","isSidechain":true,"promptId":"ae94540a-045a-4847-b16d-c2358b5c9db6","agentId":"a07e7799acd8be342","type":"user","message":{"role":"user","content":[{"tool_use_id":"toolu_01FiKmUuqesxxTt8QqgLKUYa","type":"tool_result","content":"1\timport os\n2\timport torch\n3\timport datasets\n4\timport numpy as np\n5\tfrom transformers import AutoTokenizer\n6\tfrom typing import Optional, Dict, List\n7\tfrom torch.utils.data import Sampler\n8\tfrom collections import defaultdict\n9\t\n10\t# Disable tokenizers parallelism to avoid warnings when using multiprocessing\n11\tos.environ[\"TOKENIZERS_PARALLELISM\"] = \"false\"\n12\t\n13\t\n14\t\n15\t\n16\tdef build_interaction_keys(c1_train, c2_train, c1_val, c2_val, multiplier=1_000_000):\n17\t \"\"\"\n18\t Builds a global set of interaction keys (ClusterA * M + ClusterB) \n19\t from both train and validation sets. This serves as the Ground Truth \n20\t to prevent false negatives during training.\n21\t \n22\t Args:\n23\t c1_train: seq1_cluster array for training set\n24\t c2_train: seq2_cluster array for training set\n25\t c1_val: seq1_cluster array for validation set\n26\t c2_val: seq2_cluster array for validation set\n27\t multiplier: Multiplier for hash key generation (default 1M)\n28\t \n29\t Returns:\n30\t tuple: (valid_keys_tensor, multiplier) where valid_keys_tensor is a 1D tensor of unique interaction keys\n31\t \"\"\" \n32\t # Concatenate train and val\n33\t c1 = np.concatenate([c1_train, c1_val])\n34\t c2 = np.concatenate([c2_train, c2_val])\n35\t \n36\t # 2. Filter invalid clusters (-1)\n37\t valid_mask = (c1 != -1) & (c2 != -1)\n38\t c1 = c1[valid_mask]\n39\t c2 = c2[valid_mask]\n40\t \n41\t # 3. Validate cluster IDs don't exceed multiplier (prevent hash collisions)\n42\t max_cluster = max(c1.max(), c2.max()) if len(c1) > 0 else 0\n43\t assert max_cluster < multiplier, (\n44\t f\"Cluster ID {max_cluster} exceeds multiplier {multiplier}. \"\n45\t f\"Hash collisions will occur! Increase multiplier or reduce cluster IDs.\"\n46\t )\n47\t \n48\t # 4. Create Bidirectional Keys (Vectorized)\n49\t # Key = ClusterA * Multiplier + ClusterB\n50\t keys_fwd = c1 * multiplier + c2\n51\t keys_rev = c2 * multiplier + c1 # Symmetric: If A binds B, B binds A\n52\t \n53\t # 5. Concatenate and Unique\n54\t if len(keys_fwd) == 0:\n55\t print(\"WARNING: No valid cluster pairs found. All entries have cluster_id=-1.\")\n56\t print(\"False negative masking will be disabled.\")\n57\t return torch.empty(0, dtype=torch.long), multiplier\n58\t \n59\t unique_keys = np.unique(np.concatenate([keys_fwd, keys_rev]))\n60\t \n61\t # Return as Tensor and multiplier\n62\t return torch.from_numpy(unique_keys).long(), multiplier\n63\t\n64\t\n65\tclass ClusterSampler(Sampler):\n66\t \"\"\"\n67\t Sampler that draws from multiple dataset pools with explicit per-dataset probabilities.\n68\t\n69\t Pool 0 (PPI): cluster-weighted — sample a cluster uniformly, then an example from it.\n70\t Pools 1..N (DDI, AFDB, ...): uniform — passed in as explicit index arrays.\n71\t\n72\t sampling_probs[i] is the probability of drawing from pool i. Must sum to 1.\n73\t Compatible with Accelerate/DDP via set_epoch().\n74\t \"\"\"\n75\t def __init__(\n76\t self,\n77\t c1: np.ndarray,\n78\t c2: np.ndarray,\n79\t extra_pools: list = None, # list of np.ndarray, one per extra dataset\n80\t sampling_probs: list = None, # [ppi_prob, pool0_prob, pool1_prob, ...]\n81\t seed: int = 42,\n82\t ):\n83\t super().__init__()\n84\t self.num_samples = len(c1)\n85\t self.seed = seed\n86\t self.epoch = 0\n87\t\n88\t # Build cluster-weighted index for PPI examples (c1[i] >= 0 and c2[i] >= 0)\n89\t cluster_to_indices = defaultdict(set)\n90\t for idx in range(len(c1)):\n91\t if c1[idx] >= 0 and c2[idx] >= 0:\n92\t cluster_to_indices[int(c1[idx])].add(idx)\n93\t cluster_to_indices[int(c2[idx])].add(idx)\n94\t self.clusters = list(cluster_to_indices.keys())\n95\t self.cluster_to_indices = {k: np.array(list(v), dtype=np.int64) for k, v in cluster_to_indices.items()}\n96\t\n97\t # Extra pools (DDI, AFDB, ...)\n98\t self.extra_pools = extra_pools or []\n99\t\n100\t # Sampling probabilities: [ppi_prob, pool0_prob, pool1_prob, ...]\n101\t n_pools = 1 + len(self.extra_pools)\n102\t if sampling_probs is not None:\n103\t assert len(sampling_probs) == n_pools, \\\n104\t f\"sampling_probs length {len(sampling_probs)} must match number of pools {n_pools}\"\n105\t assert abs(sum(sampling_probs) - 1.0) < 1e-6, \"sampling_probs must sum to 1.0\"\n106\t self.sampling_probs = sampling_probs\n107\t else:\n108\t self.sampling_probs = [1.0 / n_pools] * n_pools\n109\t\n110\t # Cumulative probs for O(1) pool selection; clamp last entry to 1.0 to\n111\t # guard against floating-point underflow causing missed selection.\n112\t self.cum_probs = []\n113\t cumsum = 0.0\n114\t for p in self.sampling_probs:\n115\t cumsum += p\n116\t self.cum_probs.append(cumsum)\n117\t self.cum_probs[-1] = 1.0\n118\t\n119\t self.has_clusters = len(self.clusters) > 0\n120\t\n121\t def __iter__(self):\n122\t g = torch.Generator()\n123\t g.manual_seed(self.seed + self.epoch)\n124\t\n125\t indices = []\n126\t for _ in range(self.num_samples):\n127\t r = torch.rand(1, generator=g).item()\n128\t\n129\t # Select pool; default to last pool so floating-point edge cases\n130\t # (r >= all cum_probs) never silently fall through to PPI.\n131\t pool_idx = len(self.cum_probs) - 1\n132\t for i, cp in enumerate(self.cum_probs):\n133\t if r < cp:\n134\t pool_idx = i\n135\t break\n136\t\n137\t if pool_idx == 0:\n138\t # PPI: cluster-weighted sampling\n139\t if self.has_clusters:\n140\t cluster_idx = torch.randint(len(self.clusters), (1,), generator=g).item()\n141\t cluster = self.clusters[cluster_idx]\n142\t cluster_examples = self.cluster_to_indices[cluster]\n143\t idx = cluster_examples[torch.randint(len(cluster_examples), (1,), generator=g).item()]\n144\t else:\n145\t idx = torch.randint(self.num_samples, (1,), generator=g).item()\n146\t else:\n147\t # Extra pool: uniform sampling\n148\t pool = self.extra_pools[pool_idx - 1]\n149\t idx = pool[torch.randint(len(pool), (1,), generator=g).item()]\n150\t\n151\t indices.append(int(idx))\n152\t\n153\t return iter(indices)\n154\t\n155\t def __len__(self):\n156\t return self.num_samples\n157\t\n158\t def set_epoch(self, epoch: int):\n159\t self.epoch = epoch\n160\t\n161\t\n162\tdef create_contact_targets(\n163\t seq1_lens: List[int],\n164\t seq2_lens: List[int],\n165\t contacts_list: List[List],\n166\t target_map_size: int\n167\t) -> torch.Tensor:\n168\t \"\"\"\n169\t Create contact target tensor from list of contacts.\n170\t \n171\t Args:\n172\t seq1_lens: List of sequence 1 lengths\n173\t seq2_lens: List of sequence 2 lengths\n174\t contacts_list: List of contact lists, each contact is [i, j]\n175\t target_map_size: Size of the target map\n176\t \n177\t Returns:\n178\t contact_targets: (batch_size, target_map_size, target_map_size) bool tensor\n179\t \"\"\"\n180\t batch_size = len(seq1_lens)\n181\t contact_targets = torch.zeros(batch_size, target_map_size, target_map_size, dtype=torch.bool)\n182\t \n183\t batch_indices_list = []\n184\t r1_indices_list = []\n185\t r2_indices_list = []\n186\t \n187\t for i, (seq1_len, seq2_len, contacts) in enumerate(zip(seq1_lens, seq2_lens, contacts_list)):\n188\t seq1_len = min(seq1_len, target_map_size)\n189\t seq2_len = min(seq2_len, target_map_size)\n190\t \n191\t if contacts:\n192\t contacts_tensor = torch.tensor(contacts, dtype=torch.long)\n193\t valid_mask = (contacts_tensor[:, 0] < seq1_len) & (contacts_tensor[:, 1] < seq2_len)\n194\t valid_contacts = contacts_tensor[valid_mask]\n195\t \n196\t if len(valid_contacts) > 0:\n197\t batch_indices_list.append(torch.full((len(valid_contacts),), i, dtype=torch.long))\n198\t r1_indices_list.append(valid_contacts[:, 0])\n199\t r2_indices_list.append(valid_contacts[:, 1])\n200\t \n201\t if batch_indices_list:\n202\t batch_indices = torch.cat(batch_indices_list)\n203\t r1_indices = torch.cat(r1_indices_list)\n204\t r2_indices = torch.cat(r2_indices_list)\n205\t contact_targets[batch_indices, r1_indices, r2_indices] = True\n206\t \n207\t return contact_targets\n208\t\n209\t\n210\tclass DataCollatorForPPI:\n211\t \"\"\"\n212\t Custom data collator.\n213\t - Tokenizes seq1 and seq2 separately.\n214\t - Creates the 2D contact map target tensor.\n215\t - Optionally swaps query/target with 50% probability.\n216\t \"\"\"\n217\t def __init__(self, tokenizer: AutoTokenizer, max_len: int, swap_augment: bool = True, adds_special_tokens: bool = True):\n218\t self.tokenizer = tokenizer\n219\t self.max_len = max_len\n220\t self.swap_augment = swap_augment\n221\t self.adds_special_tokens = adds_special_tokens # ESM adds CLS/SEP, gLM2 does not\n222\t\n223\t def __call__(self, batch: list[dict]):\n224\t # Randomly swap query/target for augmentation\n225\t if self.swap_augment:\n226\t batch = [\n227\t {'seq1': item['seq2'],\n228\t 'seq2': item['seq1'], \n229\t 'contacts': [(j, i) for i, j in item.get('contacts', [])],\n230\t 'seq1_cluster': item.get('seq2_cluster'),\n231\t 'seq2_cluster': item.get('seq1_cluster'),\n232\t 'complex_id': item.get('complex_id')}\n233\t if torch.rand(1).item() < 0.5 else item\n234\t for item in batch\n235\t ]\n236\t \n237\t # Extract sequences\n238\t seqs1 = [item['seq1'] for item in batch]\n239\t seqs2 = [item['seq2'] for item in batch]\n240\t \n241\t # Normalize cluster IDs -1 (use -1 for missing/DDI dataset)\n242\t seq1_clusters = [item.get('seq1_cluster') for item in batch]\n243\t seq1_clusters = [-1 if val is None else val for val in seq1_clusters]\n244\t seq2_clusters = [item.get('seq2_cluster') for item in batch]\n245\t seq2_clusters = [-1 if val is None else val for val in seq2_clusters]\n246\t \n247\t # Tokenize separately (ESM adds CLS/SEP, gLM2 does not)\n248\t batch1_tokens = self.tokenizer(\n249\t seqs1, \n250\t return_tensors='pt', \n251\t padding='max_length', \n252\t truncation=True, \n253\t max_length=self.max_len\n254\t )\n255\t \n256\t batch2_tokens = self.tokenizer(\n257\t seqs2, \n258\t return_tensors='pt', \n259\t padding='max_length', \n260\t truncation=True, \n261\t max_length=self.max_len\n262\t )\n263\t \n264\t # Create contact targets using standalone function\n265\t seq1_lens = [len(item['seq1']) for item in batch]\n266\t seq2_lens = [len(item['seq2']) for item in batch]\n267\t # The target map size matches model output (which slices off CLS/SEP if present)\n268\t target_map_size = self.max_len - 2 if self.adds_special_tokens else self.max_len\n269\t contacts_list = [item.get('contacts', []) for item in batch]\n270\t contact_targets = create_contact_targets(seq1_lens, seq2_lens, contacts_list, target_map_size)\n271\t\n272\t # Carry is_positive flag if present (defaults to True for backward compatibility/PDB_PPI)\n273\t is_positive = [item.get('is_positive', True) for item in batch]\n274\t\n275\t return {\n276\t 'seq1_tokens': batch1_tokens,\n277\t 'seq2_tokens': batch2_tokens,\n278\t 'contact_target': contact_targets,\n279\t 'seq1_cluster': torch.tensor(seq1_clusters, dtype=torch.long),\n280\t 'seq2_cluster': torch.tensor(seq2_clusters, dtype=torch.long),\n281\t 'is_positive': torch.tensor(is_positive, dtype=torch.bool)\n282\t }\n283\t\n284\t\n285\tclass BenchmarkCollator:\n286\t \"\"\"\n287\t Simple data collator for benchmark evaluation.\n288\t Tokenizes seq1 and seq2 separately and includes labels.\n289\t \"\"\"\n290\t def __init__(self, tokenizer: AutoTokenizer, max_len: int, adds_special_tokens: bool = True):\n291\t self.tokenizer = tokenizer\n292\t self.max_len = max_len\n293\t self.adds_special_tokens = adds_special_tokens\n294\t \n295\t def __call__(self, batch: list[dict]):\n296\t seqs1 = [item['seq1'] for item in batch]\n297\t seqs2 = [item['seq2'] for item in batch]\n298\t labels = torch.tensor([item['label'] for item in batch], dtype=torch.long)\n299\t \n300\t batch1_tokens = self.tokenizer(\n301\t seqs1, return_tensors='pt', padding='max_length', \n302\t truncation=True, max_length=self.max_len\n303\t )\n304\t batch2_tokens = self.tokenizer(\n305\t seqs2, return_tensors='pt', padding='max_length', \n306\t truncation=True, max_length=self.max_len\n307\t )\n308\t \n309\t return {\n310\t 'seq1_tokens': {'input_ids': batch1_tokens['input_ids'], 'attention_mask': batch1_tokens['attention_mask']},\n311\t 'seq2_tokens': {'input_ids': batch2_tokens['input_ids'], 'attention_mask': batch2_tokens['attention_mask']},\n312\t 'label': labels\n313\t }\n314\t\n315\t\n316\tdef load_retrieval_dataset(\n317\t dataset_path: str,\n318\t tokenizer: AutoTokenizer,\n319\t max_len: int\n320\t) -> Optional[Dict]:\n321\t \"\"\"\n322\t Load and tokenize retrieval evaluation dataset.\n323\t \n324\t Args:\n325\t dataset_path: Path to retrieval dataset (HF hub or local)\n326\t tokenizer: Tokenizer to use\n327\t max_len: Maximum sequence length\n328\t \n329\t Returns:\n330\t Dict with 'queries', 'corpus', 'query_inputs', 'corpus_inputs' or None if loading fails\n331\t \"\"\"\n332\t retrieval_ds = datasets.load_from_disk(dataset_path)\n333\t\n334\t \n335\t queries = list(retrieval_ds['queries'])\n336\t corpus = list(retrieval_ds['corpus'])\n337\t \n338\t # Tokenize sequences\n339\t query_sequences = [q['sequence'] for q in queries]\n340\t corpus_sequences = [c['sequence'] for c in corpus]\n341\t \n342\t query_inputs = tokenizer(\n343\t query_sequences,\n344\t return_tensors='pt',\n345\t padding='max_length',\n346\t truncation=True,\n347\t max_length=max_len\n348\t )\n349\t corpus_inputs = tokenizer(\n350\t corpus_sequences,\n351\t return_tensors='pt',\n352\t padding='max_length',\n353\t truncation=True,\n354\t max_length=max_len\n355\t )\n356\t \n357\t return {\n358\t 'queries': queries,\n359\t 'corpus': corpus,\n360\t 'query_inputs': query_inputs,\n361\t 'corpus_inputs': corpus_inputs\n362\t }\n363\t"}]},"uuid":"75948e45-fde4-44c0-a3d5-e54d8f40577b","timestamp":"2026-06-01T15:28:21.306Z","sourceToolAssistantUUID":"b3bf75ea-639e-4193-9cf7-f47b8a9939c0","userType":"external","entrypoint":"claude-vscode","cwd":"/home/ubuntu/ppi_project","sessionId":"5a87a0cd-6802-4ac5-a708-29b2bf7ea563","version":"2.1.159","gitBranch":"afdb"} +{"parentUuid":"b88a253a-19a3-42d0-932f-4e91525e7939","isSidechain":true,"promptId":"ae94540a-045a-4847-b16d-c2358b5c9db6","agentId":"a07e7799acd8be342","type":"user","message":{"role":"user","content":[{"tool_use_id":"toolu_016aL1jZdYhhjwF8otCoAhR2","type":"tool_result","content":"1\timport torch\n2\timport torch.nn as nn\n3\timport torch.nn.functional as F\n4\tfrom typing import Literal, Optional\n5\tfrom transformers import AutoModel\n6\tfrom config import ModelConfig\n7\tfrom layers import PairformerConfig, MLPHead, TransformerLayers\n8\t\n9\t\n10\tdef load_plm_model(model_name: str):\n11\t \"\"\"\n12\t Load a protein language model (PLM).\n13\t \n14\t If model_name contains \"esm\", loads using FAEsmForMaskedLM from faesm.esm.\n15\t Otherwise, loads using AutoModel.from_pretrained.\n16\t \n17\t Args:\n18\t model_name (str): Name/path of the model to load\n19\t \n20\t Returns:\n21\t nn.Module: Loaded PLM model\n22\t \"\"\"\n23\t if \"esm\" in model_name.lower():\n24\t try:\n25\t from plm_models.modeling_esm_flash import EsmFlashModel\n26\t esm_model = EsmFlashModel.from_pretrained(model_name)\n27\t except ImportError:\n28\t print(\"WARNING: Could not load ESM model with flash attention. Falling back to AutoModel.\")\n29\t esm_model = AutoModel.from_pretrained(model_name, trust_remote_code=True)\n30\t # Delete some unused modules to make DDP work\n31\t esm_model.pooler = None\n32\t esm_model.embeddings.position_embeddings = None\n33\t return esm_model\n34\t elif \"glm2\" in model_name.lower():\n35\t try:\n36\t from plm_models.modeling_glm2_flash import gLM2Model\n37\t return gLM2Model.from_pretrained(model_name)\n38\t except ImportError:\n39\t print(\"WARNING: Could not load gLM2 model with flash attention. Falling back to AutoModel.\")\n40\t return AutoModel.from_pretrained(model_name, trust_remote_code=True)\n41\t else:\n42\t return AutoModel.from_pretrained(model_name, trust_remote_code=True)\n43\t\n44\t\n45\tdef freeze_plm_layers(model: nn.Module, num_layers: int):\n46\t \"\"\"\n47\t Freeze the first num_layers of a transformer model.\n48\t \n49\t Supports various transformer architectures:\n50\t - ESM models: encoder.layer\n51\t - gLM2 models: encoder.layers\n52\t - Standard transformers: encoder.layer or layers\n53\t \n54\t Args:\n55\t model (nn.Module): The PLM model\n56\t num_layers (int): Number of layers to freeze (from the beginning)\n57\t \"\"\"\n58\t if num_layers <= 0:\n59\t return\n60\t \n61\t # Try to find the encoder layers\n62\t layers = None\n63\t if hasattr(model, 'encoder') and hasattr(model.encoder, 'layer'):\n64\t layers = model.encoder.layer\n65\t elif hasattr(model, 'encoder') and hasattr(model.encoder, 'layers'):\n66\t layers = model.encoder.layers\n67\t elif hasattr(model, 'layers'):\n68\t layers = model.layers\n69\t \n70\t if layers is None:\n71\t print(f\"WARNING: Could not find encoder layers in model. Skipping layer freezing.\")\n72\t return\n73\t \n74\t num_frozen = 0\n75\t for layer in layers[:num_layers]:\n76\t layer.requires_grad_(False)\n77\t num_frozen += 1\n78\t\n79\t # Freeze embeddings\n80\t if hasattr(model, \"embeddings\"):\n81\t model.embeddings.requires_grad_(False)\n82\t if hasattr(model, \"tok_embeddings\"):\n83\t model.tok_embeddings.requires_grad_(False)\n84\t \n85\t print(f\"Frozen {num_frozen} out of {len(layers)} PLM layers.\")\n86\t\n87\t\n88\tdef get_contact_negatives(\n89\t embed1: torch.Tensor, \n90\t embed2: torch.Tensor, \n91\t mask1: torch.Tensor, \n92\t mask2: torch.Tensor, \n93\t random_neg_ratio: float = 1.0,\n94\t self_neg_ratio: float = 1.0,\n95\t hard_neg_ratio: float = 0.0,\n96\t clusters1: Optional[torch.Tensor] = None,\n97\t clusters2: Optional[torch.Tensor] = None,\n98\t valid_keys: Optional[torch.Tensor] = None,\n99\t multiplier: int = 1_000_000,\n100\t clip_embed1: Optional[torch.Tensor] = None,\n101\t clip_embed2: Optional[torch.Tensor] = None,\n102\t):\n103\t \"\"\"\n104\t Generate negative pairs for contact prediction, avoiding known interactions.\n105\t Supports random, self, and hard negatives (based on CLIP similarity).\n106\t \"\"\"\n107\t B = embed1.size(0)\n108\t device = embed1.device\n109\t \n110\t # Pre-compute known interaction mask (B, B) - True means it's a known interaction\n111\t known_mask = None\n112\t if clusters1 is not None and clusters2 is not None and valid_keys is not None and len(valid_keys) > 0:\n113\t batch_keys = clusters1.unsqueeze(1) * multiplier + clusters2.unsqueeze(0)\n114\t invalid = (clusters1.unsqueeze(1) == -1) | (clusters2.unsqueeze(0) == -1)\n115\t known_mask = torch.isin(batch_keys, valid_keys.to(device)) & ~invalid\n116\t \n117\t # Collect negative pair indices\n118\t neg_idx1_list, neg_idx2_list = [], []\n119\t \n120\t # Random negatives: shift pairs (A_i, B_{i+shift}), skip known interactions\n121\t if random_neg_ratio > 0 and B > 1:\n122\t num_rand = int(B * random_neg_ratio)\n123\t reps = (num_rand + B - 1) // B\n124\t for shift in range(1, reps + 1):\n125\t n = min(B, num_rand - (shift - 1) * B)\n126\t idx1 = torch.arange(n, device=device)\n127\t idx2 = (idx1 + shift) % B\n128\t if known_mask is not None:\n129\t valid = ~known_mask[idx1, idx2]\n130\t idx1, idx2 = idx1[valid], idx2[valid]\n131\t neg_idx1_list.append(idx1)\n132\t neg_idx2_list.append(idx2)\n133\t \n134\t # Self negatives: (A_i, A_i), skip known interactions\n135\t if self_neg_ratio > 0:\n136\t n = min(B, int(B * self_neg_ratio))\n137\t idx = torch.arange(n, device=device)\n138\t if known_mask is not None:\n139\t valid = ~known_mask[idx, idx]\n140\t idx = idx[valid]\n141\t neg_idx1_list.append(idx)\n142\t neg_idx2_list.append(idx)\n143\t \n144\t # Hard negatives: select most similar non-known pairs by CLIP score\n145\t if hard_neg_ratio > 0 and clip_embed1 is not None and clip_embed2 is not None and B > 1:\n146\t num_hard_per_example = max(1, int(hard_neg_ratio))\n147\t with torch.no_grad():\n148\t sim_matrix = clip_embed1 @ clip_embed2.T\n149\t # Mask out diagonal (positives)\n150\t sim_matrix.fill_diagonal_(-1e9)\n151\t # Mask out known interactions before selecting top-K\n152\t if known_mask is not None:\n153\t sim_matrix = sim_matrix.masked_fill(known_mask, -1e9)\n154\t # Get top-K most similar valid pairs per example\n155\t num_hard_actual = min(num_hard_per_example, B - 1)\n156\t _, top_indices_per_row = torch.topk(sim_matrix, num_hard_actual, dim=1)\n157\t idx1 = torch.arange(B, device=device).unsqueeze(1).expand(-1, num_hard_actual).flatten()\n158\t idx2 = top_indices_per_row.flatten()\n159\t neg_idx1_list.append(idx1)\n160\t neg_idx2_list.append(idx2)\n161\t \n162\t # Combine all negatives\n163\t if neg_idx1_list:\n164\t neg_idx1 = torch.cat(neg_idx1_list)\n165\t neg_idx2 = torch.cat(neg_idx2_list)\n166\t \n167\t if len(neg_idx1) > 0:\n168\t return embed1[neg_idx1], embed2[neg_idx2], mask1[neg_idx1], mask2[neg_idx2]\n169\t \n170\t # Return empty if no negatives\n171\t return (torch.empty(0, *embed1.shape[1:], device=device),\n172\t torch.empty(0, *embed2.shape[1:], device=device),\n173\t torch.empty(0, *mask1.shape[1:], device=device),\n174\t torch.empty(0, *mask2.shape[1:], device=device))\n175\t\n176\t\n177\tclass ContactHead(nn.Module):\n178\t \"\"\"\n179\t Computes interactions between two proteins using Multi-Head Attention logic.\n180\t Returns the raw attention logits (the contact map) for supervision.\n181\t \"\"\"\n182\t def __init__(self, input_dim, contact_dim, num_heads=8):\n183\t super().__init__()\n184\t self.num_heads = num_heads\n185\t self.head_dim = contact_dim // num_heads\n186\t assert contact_dim % num_heads == 0, \"Input dim must be divisible by num_heads\"\n187\t\n188\t\n189\t # Segment Embedding\n190\t self.segment_embed = nn.Embedding(2, input_dim)\n191\t nn.init.normal_(self.segment_embed.weight, std=0.02)\n192\t \n193\t self.transformer = TransformerLayers(PairformerConfig(dim=input_dim, heads=num_heads, depth=2, out_dim=input_dim))\n194\t\n195\t # Layer normalization before projection (stabilizes training)\n196\t self.norm = nn.LayerNorm(input_dim)\n197\t\n198\t # Projections for Query (Seq1) and Key (Seq2)\n199\t self.q_proj = nn.Linear(input_dim, contact_dim)\n200\t self.k_proj = nn.Linear(input_dim, contact_dim)\n201\t \n202\t self.output_mix = nn.Linear(num_heads, 1)\n203\t nn.init.constant_(self.output_mix.bias, -3.0)\n204\t\n205\t # Scale factor (1 / sqrt(head_dim))\n206\t self.scale = self.head_dim ** -0.5\n207\t\n208\t def forward(self, embed1, embed2, mask1, mask2):\n209\t B, L1, D = embed1.shape\n210\t _, L2, _ = embed2.shape\n211\t\n212\t seg1 = self.segment_embed(torch.zeros(L1, device=embed1.device, dtype=torch.long))\n213\t seg2 = self.segment_embed(torch.ones(L2, device=embed1.device, dtype=torch.long))\n214\t \n215\t # Concat inputs\n216\t x = torch.cat([embed1 + seg1.unsqueeze(0), embed2 + seg2.unsqueeze(0)], dim=1)\n217\t \n218\t combined_mask = None\n219\t if mask1 is not None and mask2 is not None:\n220\t combined_mask = torch.cat([mask1, mask2], dim=1).bool()\n221\t \n222\t # Run Transformer\n223\t x = self.transformer(x, attention_mask=combined_mask)\n224\t \n225\t # Split & Bottleneck\n226\t embed1 = x[:, :L1, :]\n227\t embed2 = x[:, L1:, :]\n228\t\n229\t # Normalize embeddings before projection (stabilizes training)\n230\t embed1 = self.norm(embed1)\n231\t embed2 = self.norm(embed2)\n232\t\n233\t # 1. Project and Reshape to (Batch, Heads, SeqLen, HeadDim)\n234\t # We separate heads to allow the model to learn different contact types\n235\t q = self.q_proj(embed1).view(B, L1, self.num_heads, self.head_dim).transpose(1, 2)\n236\t k = self.k_proj(embed2).view(B, L2, self.num_heads, self.head_dim).transpose(1, 2)\n237\t\n238\t # 2. Compute Attention Scores (The Contact Map)\n239\t # Shape: (Batch, Heads, L1, L2)\n240\t attn_logits = torch.matmul(q, k.transpose(-2, -1)) * self.scale\n241\t\n242\t # 3. Mix Heads (More efficient: ensure contiguous memory layout)\n243\t # Permute to (Batch, L1, L2, Heads) to apply the mixing Linear layer\n244\t attn_logits = attn_logits.permute(0, 2, 3, 1).contiguous()\n245\t \n246\t # Project heads -> 1 final score per pair\n247\t # Shape: (Batch, L1, L2, 1) -> (Batch, L1, L2)\n248\t contact_logits = self.output_mix(attn_logits).squeeze(-1)\n249\t\n250\t if mask1 is not None and mask2 is not None:\n251\t # Shape: (Batch, L1, L2)\n252\t valid_mask = (mask1.unsqueeze(2) * mask2.unsqueeze(1)).bool()\n253\t else:\n254\t valid_mask = torch.ones_like(contact_logits, dtype=torch.bool)\n255\t\n256\t return contact_logits, valid_mask\n257\t\n258\t\n259\tclass PoolLayer(nn.Module):\n260\t def __init__(self, method: Literal[\"mean\", \"max\"] = \"mean\") -> None:\n261\t super().__init__()\n262\t self.method = method\n263\t \n264\t def forward(self, embeds: torch.FloatTensor, mask: torch.BoolTensor) -> torch.FloatTensor:\n265\t mask = mask.unsqueeze(-1).bool()\n266\t if self.method == \"mean\":\n267\t embeds = torch.where(mask, embeds, 0.0)\n268\t embeds = torch.sum(embeds, -2)\n269\t embeds /= torch.clamp(torch.sum(mask, dim=-2, dtype=embeds.dtype), min=1.0)\n270\t elif self.method == \"max\":\n271\t embeds = torch.where(mask, embeds, torch.finfo(embeds.dtype).min)\n272\t embeds = torch.max(embeds, dim=-2).values\n273\t embeds = torch.where(mask.any(-2), embeds, torch.zeros_like(embeds))\n274\t return embeds\n275\t\n276\tclass ContrastiveHead(nn.Module):\n277\t \"\"\"\n278\t Linear projection, pooling, and MLP head, normalization for contrastive learning.\n279\t \"\"\"\n280\t def __init__(self, hidden_dim, clip_embed_dim):\n281\t super().__init__()\n282\t self.pool_layer = PoolLayer(method=\"mean\")\n283\t self.head = MLPHead(hidden_dim, clip_embed_dim)\n284\t\n285\t def forward(self, residue_embeds, mask):\n286\t pool_embeds = self.pool_layer(residue_embeds, mask)\n287\t return F.normalize(self.head(pool_embeds), dim=-1)\n288\t\n289\tclass ContrastivePPIModel(nn.Module):\n290\t def __init__(\n291\t self, \n292\t config: ModelConfig, \n293\t enable_contact_prediction: bool = True, \n294\t swap_self_negative: bool = False, \n295\t enable_clip: bool = True,\n296\t contact_random_neg_ratio: float = 1.0,\n297\t contact_self_neg_ratio: float = 1.0,\n298\t contact_hard_neg_ratio: float = 0.0,\n299\t valid_keys_tensor: Optional[torch.Tensor] = None,\n300\t multiplier: int = 1_000_000,\n301\t ):\n302\t super().__init__()\n303\t self.config = config\n304\t self.enable_contact_prediction = enable_contact_prediction\n305\t self.swap_self_negative = swap_self_negative and enable_clip\n306\t self.enable_clip = enable_clip\n307\t self.contact_random_neg_ratio = contact_random_neg_ratio\n308\t self.contact_self_neg_ratio = contact_self_neg_ratio\n309\t self.contact_hard_neg_ratio = contact_hard_neg_ratio\n310\t # Register valid_keys as buffer for contact negative filtering\n311\t if valid_keys_tensor is not None:\n312\t self.multiplier = multiplier\n313\t self.register_buffer('valid_keys', valid_keys_tensor, persistent=False)\n314\t\n315\t \n316\t # 1. SHARED Backbone\n317\t self.plm = load_plm_model(config.plm_model_name)\n318\t \n319\t # Freeze logic\n320\t if config.num_frozen_layers > 0:\n321\t freeze_plm_layers(self.plm, config.num_frozen_layers)\n322\t\n323\t if hasattr(self.plm.config, 'hidden_size'):\n324\t hidden_dim = self.plm.config.hidden_size\n325\t else:\n326\t hidden_dim = self.plm.config.dim\n327\t \n328\t # ESM tokenizer adds cls/eos tokens.\n329\t self._adds_special_tokens = \"esm\" in config.plm_model_name.lower()\n330\t if self.enable_clip:\n331\t self.head_q = ContrastiveHead(hidden_dim, config.clip_embed_dim)\n332\t self.head_k = ContrastiveHead(hidden_dim, config.clip_embed_dim)\n333\t self.logit_scale = nn.Parameter(torch.ones([]) * torch.log(torch.tensor(1 / 0.07)))\n334\t else:\n335\t self.head_q = None\n336\t self.head_k = None\n337\t self.logit_scale = torch.tensor(1.0)\n338\t\n339\t \n340\t \n341\t # Contact Head\n342\t self.contact_head = ContactHead(hidden_dim, config.contact_embed_dim, num_heads=8)\n343\t # self.contact_head = PairformerContactHead(PairformerConfig(dim=hidden_dim, heads=8, depth=3, out_dim=config.contact_embed_dim))\n344\t\n345\t def gradient_checkpointing_enable(self):\n346\t \"\"\"Enable gradient checkpointing for the PLM backbone and contact head transformer if supported.\"\"\"\n347\t if hasattr(self.plm, 'gradient_checkpointing_enable'):\n348\t self.plm.gradient_checkpointing_enable()\n349\t self.contact_head.transformer.gradient_checkpointing = True\n350\t \n351\t def gradient_checkpointing_disable(self):\n352\t \"\"\"Disable gradient checkpointing for the PLM backbone and contact head transformer if supported.\"\"\"\n353\t if hasattr(self.plm, 'gradient_checkpointing_disable'):\n354\t self.plm.gradient_checkpointing_disable()\n355\t self.contact_head.transformer.gradient_checkpointing = False\n356\t\n357\t\n358\t def encode_protein(\n359\t self,\n360\t input_ids: torch.Tensor,\n361\t attention_mask: Optional[torch.Tensor] = None,\n362\t ):\n363\t \"\"\"\n364\t Runs the shared plm model on tokenized protein sequences.\n365\t \n366\t Args:\n367\t input_ids (torch.Tensor): (B, L) - Token IDs from the tokenizer\n368\t attention_mask (Optional[torch.Tensor]): (B, L) - Attention mask (1 for real tokens, 0 for padding)\n369\t \n370\t Returns:\n371\t residue_embeds (torch.Tensor): (B, L, D_plm) - Embeddings for each residue\n372\t pool_embed (torch.Tensor): (B, D_plm) - Pooled embedding using mask-aware pooling\n373\t \"\"\"\n374\t outputs = self.plm(input_ids=input_ids, attention_mask=attention_mask)\n375\t \n376\t # (B, L, D_plm) - Full sequence embeddings\n377\t residue_embeds = outputs.last_hidden_state\n378\t return residue_embeds\n379\t\n380\t def predict_contacts(self, embed1: torch.Tensor, embed2: torch.Tensor, mask1: torch.Tensor, mask2: torch.Tensor):\n381\t \"\"\"\n382\t Predicts a contact map using dot product between projected residue embeddings.\n383\t \n384\t Args:\n385\t embed1 (torch.Tensor): (B, L1, D_plm) - Residue embeddings for sequence 1 (Query)\n386\t embed2 (torch.Tensor): (B, L2, D_plm) - Residue embeddings for sequence 2 (Key)\n387\t mask1 (torch.Tensor): (B, L1) - Attention mask for seq1\n388\t mask2 (torch.Tensor): (B, L2) - Attention mask for seq2\n389\t \n390\t Returns:\n391\t contact_logits (torch.Tensor): (B, L1, L2) - Contact prediction logits\n392\t \"\"\"\n393\t return self.contact_head(embed1, embed2, mask1, mask2)\n394\t \n395\t def encode_pair(\n396\t self,\n397\t seq1_input_ids: torch.Tensor,\n398\t seq2_input_ids: torch.Tensor,\n399\t seq1_attention_mask: Optional[torch.Tensor] = None,\n400\t seq2_attention_mask: Optional[torch.Tensor] = None,\n401\t ):\n402\t \"\"\"\n403\t Encodes a pair of protein sequences through PLM in a single batched call.\n404\t Returns PLM-level representations (residue embeddings and pooled embeddings).\n405\t \"\"\"\n406\t B = seq1_input_ids.shape[0]\n407\t \n408\t # Batch both sequences together for a single PLM call\n409\t # Concatenate along batch dimension: (2*B, L)\n410\t batched_input_ids = torch.cat([seq1_input_ids, seq2_input_ids], dim=0)\n411\t if seq1_attention_mask is not None and seq2_attention_mask is not None:\n412\t batched_attention_mask = torch.cat([seq1_attention_mask, seq2_attention_mask], dim=0)\n413\t else:\n414\t batched_attention_mask = None\n415\t \n416\t # Single PLM call: (2*B, L, D_plm)\n417\t batched_residue_embeds= self.encode_protein(batched_input_ids, batched_attention_mask)\n418\t \n419\t # Split back into seq1 and seq2\n420\t residue_embed1 = batched_residue_embeds[:B] # (B, L1, D_plm)\n421\t residue_embed2 = batched_residue_embeds[B:] # (B, L2, D_plm)\n422\t \n423\t return residue_embed1, residue_embed2\n424\t \n425\t def embed_pair(\n426\t self,\n427\t seq1_input_ids: torch.Tensor,\n428\t seq2_input_ids: torch.Tensor,\n429\t seq1_attention_mask: Optional[torch.Tensor] = None,\n430\t seq2_attention_mask: Optional[torch.Tensor] = None,\n431\t seq1_cluster: Optional[torch.Tensor] = None,\n432\t seq2_cluster: Optional[torch.Tensor] = None,\n433\t ):\n434\t \"\"\"\n435\t Computes standard embeddings for CLIP loss, AND 'self' embeddings for identity penalty.\n436\t No batch doubling.\n437\t \n438\t Returns:\n439\t q1 (torch.Tensor): (B, D_clip) - Q(A) - Normalized query embedding for seq1\n440\t k2 (torch.Tensor): (B, D_clip) - K(B) - Normalized key embedding for seq2\n441\t k1 (torch.Tensor): (B, D_clip) - K(A) - Normalized key embedding for seq1 (for identity check)\n442\t q2 (torch.Tensor): (B, D_clip) - Q(B) - Normalized query embedding for seq2 (for identity check)\n443\t seq1_cluster_out: Processed cluster IDs for seq1 (doubled if swap_self_negative)\n444\t seq2_cluster_out: Processed cluster IDs for seq2 (doubled if swap_self_negative)\n445\t \"\"\"\n446\t residue_embed1, residue_embed2 = self.encode_pair(\n447\t seq1_input_ids, seq2_input_ids, seq1_attention_mask, seq2_attention_mask\n448\t )\n449\t \n450\t q = self.head_q(residue_embed1, seq1_attention_mask)\n451\t k = self.head_k(residue_embed2, seq2_attention_mask)\n452\t\n453\t if self.swap_self_negative and self.training:\n454\t q2 = self.head_q(residue_embed2, seq2_attention_mask)\n455\t k1 = self.head_k(residue_embed1, seq1_attention_mask)\n456\t q = torch.cat([q, q2], dim=0)\n457\t k = torch.cat([k, k1], dim=0)\n458\t \n459\t # Handle cluster ID swapping\n460\t if seq1_cluster is not None and seq2_cluster is not None:\n461\t seq1_cluster = torch.cat([seq1_cluster, seq2_cluster], dim=0)\n462\t seq2_cluster = torch.cat([seq2_cluster, seq1_cluster[:seq1_cluster.shape[0]//2]], dim=0)\n463\t\n464\t return q, k, seq1_cluster, seq2_cluster\n465\t\n466\t def forward(\n467\t self,\n468\t seq1_input_ids: torch.Tensor,\n469\t seq2_input_ids: torch.Tensor,\n470\t seq1_attention_mask: Optional[torch.Tensor] = None,\n471\t seq2_attention_mask: Optional[torch.Tensor] = None,\n472\t seq1_cluster: Optional[torch.Tensor] = None,\n473\t seq2_cluster: Optional[torch.Tensor] = None,\n474\t ):\n475\t \"\"\"\n476\t Full forward pass for both CLIP and contact prediction.\n477\t \"\"\"\n478\t seq1_mask = seq1_attention_mask.bool()\n479\t seq2_mask = seq2_attention_mask.bool()\n480\t residue_embed1, residue_embed2 = self.encode_pair(\n481\t seq1_input_ids, seq2_input_ids, seq1_mask, seq2_mask\n482\t )\n483\t\n484\t if self.enable_clip:\n485\t q = self.head_q(residue_embed1, seq1_mask)\n486\t k = self.head_k(residue_embed2, seq2_mask)\n487\t if self.swap_self_negative and self.training:\n488\t q2 = self.head_q(residue_embed2, seq2_mask)\n489\t k1 = self.head_k(residue_embed1, seq1_mask)\n490\t q = torch.cat([q, q2], dim=0)\n491\t k = torch.cat([k, k1], dim=0)\n492\t \n493\t # Handle cluster ID swapping\n494\t if seq1_cluster is not None and seq2_cluster is not None:\n495\t orig_seq1 = seq1_cluster\n496\t orig_seq2 = seq2_cluster\n497\t seq1_cluster = torch.cat([orig_seq1, orig_seq2], dim=0)\n498\t seq2_cluster = torch.cat([orig_seq2, orig_seq1], dim=0)\n499\t else:\n500\t q = None\n501\t k = None\n502\t\n503\t # --- 2. Contact Prediction Path (Residue-level) ---\n504\t \n505\t if self.enable_contact_prediction:\n506\t B = residue_embed1.shape[0] # Original batch size\n507\t \n508\t # Store original cluster IDs before CLIP swapping for contact prediction\n509\t orig_seq1_cluster = seq1_cluster\n510\t orig_seq2_cluster = seq2_cluster\n511\t \n512\t # Slice out [CLS] and [SEP] tokens if the tokenizer added them\n513\t if self._adds_special_tokens:\n514\t residue_embed1 = residue_embed1[:, 1:-1, :]\n515\t residue_embed2 = residue_embed2[:, 1:-1, :]\n516\t seq1_mask = seq1_mask[:, 1:-1] if seq1_mask is not None else None\n517\t seq2_mask = seq2_mask[:, 1:-1] if seq2_mask is not None else None\n518\t\n519\t # Prepare embeddings for negative mining\n520\t neg_embed1, neg_embed2 = residue_embed1, residue_embed2\n521\t neg_mask1, neg_mask2 = seq1_mask, seq2_mask\n522\t neg_clusters1, neg_clusters2 = orig_seq1_cluster, orig_seq2_cluster\n523\t \n524\t if self.training:\n525\t if self.swap_self_negative:\n526\t neg_embed1 = torch.cat([residue_embed1, residue_embed2], dim=0)\n527\t neg_embed2 = torch.cat([residue_embed2, residue_embed1], dim=0)\n528\t neg_mask1 = torch.cat([seq1_mask, seq2_mask], dim=0)\n529\t neg_mask2 = torch.cat([seq2_mask, seq1_mask], dim=0)\n530\t neg_clusters1 = seq1_cluster # Already doubled\n531\t neg_clusters2 = seq2_cluster\n532\t\n533\t # Get negatives (using potentially doubled embeddings for hard mining)\n534\t neg1, neg2, neg_mask1, neg_mask2 = get_contact_negatives(\n535\t neg_embed1, neg_embed2, neg_mask1, neg_mask2,\n536\t random_neg_ratio=self.contact_random_neg_ratio,\n537\t self_neg_ratio=self.contact_self_neg_ratio,\n538\t hard_neg_ratio=self.contact_hard_neg_ratio,\n539\t clusters1=neg_clusters1,\n540\t clusters2=neg_clusters2,\n541\t valid_keys=self.valid_keys,\n542\t multiplier=self.multiplier,\n543\t clip_embed1=q if q is not None else None,\n544\t clip_embed2=k if k is not None else None,\n545\t )\n546\t \n547\t # Concatenate positives + negatives (works even if neg is empty)\n548\t residue_embed1 = torch.cat([residue_embed1, neg1], dim=0)\n549\t residue_embed2 = torch.cat([residue_embed2, neg2], dim=0)\n550\t seq1_mask = torch.cat([seq1_mask, neg_mask1], dim=0)\n551\t seq2_mask = torch.cat([seq2_mask, neg_mask2], dim=0)\n552\t \n553\t contact_logits, contact_valid_mask = self.predict_contacts(\n554\t residue_embed1, residue_embed2, seq1_mask, seq2_mask\n555\t )\n556\t else:\n557\t # Return dummy contact logits (will not be used)\n558\t contact_logits = None\n559\t contact_valid_mask = None\n560\t \n561\t return q, k, contact_logits, self.logit_scale.exp(), contact_valid_mask, seq1_cluster, seq2_cluster\n562\t\n563\t\n564\t\n565\t\n566\t\n567\t\n568\t\n569\t\n570\t# class LatentPooling(nn.Module):\n571\t# \"\"\" Learns K distinct vectors using SDPA\"\"\"\n572\t\n573\t# def __init__(self, hidden_dim, num_latents=16, output_dim=None, num_heads=4):\n574\t# super().__init__()\n575\t# self.num_latents = num_latents\n576\t# self.output_dim = output_dim or hidden_dim\n577\t# self.num_heads = num_heads\n578\t# self.head_dim = hidden_dim // num_heads\n579\t# assert hidden_dim % num_heads == 0, \"hidden_dim must be divisible by num_heads\"\n580\t\n581\t# # 1. The Learnable Latent Queries (The \"Questions\")\n582\t# self.latent_queries = nn.Parameter(torch.randn(1, num_latents, hidden_dim))\n583\t \n584\t# # 2. Projections for SDPA\n585\t# self.q_proj = nn.Linear(hidden_dim, hidden_dim)\n586\t# self.k_proj = nn.Linear(hidden_dim, hidden_dim)\n587\t# self.v_proj = nn.Linear(hidden_dim, hidden_dim)\n588\t \n589\t# # 3. Output Projection\n590\t# self.o_proj = nn.Linear(hidden_dim, self.output_dim)\n591\t\n592\t# def forward(self, seq_embeds, attention_mask=None):\n593\t# \"\"\"\n594\t# Args:\n595\t# seq_embeds: (B, L, D) - From PLM\n596\t# attention_mask: (B, L) - 1 for valid, 0 for pad\n597\t# Returns:\n598\t# latents: (B, K, OutDim)\n599\t# \"\"\"\n600\t# B, L, D = seq_embeds.shape\n601\t# K = self.num_latents\n602\t \n603\t# # --- Prepare Q, K, V ---\n604\t \n605\t# # Query: Expand latents (B, K, D)\n606\t# queries = self.latent_queries.expand(B, -1, -1)\n607\t# q = self.q_proj(queries).view(B, K, self.num_heads, self.head_dim).transpose(1, 2)\n608\t# # q: (B, Heads, K, HeadDim)\n609\t \n610\t# # Key/Value: Protein Sequence (B, L, D)\n611\t# k = self.k_proj(seq_embeds).view(B, L, self.num_heads, self.head_dim).transpose(1, 2)\n612\t# v = self.v_proj(seq_embeds).view(B, L, self.num_heads, self.head_dim).transpose(1, 2)\n613\t# # k, v: (B, Heads, L, HeadDim)\n614\t \n615\t# # --- Prepare Mask for SDPA ---\n616\t# # SDPA attn_mask: Binary mask where True indicates UNVALID (masked) positions.\n617\t# # We need to broadcast (B, L) -> (B, 1, 1, L) so it applies to all heads/queries\n618\t# attn_mask = None\n619\t# if attention_mask is not None:\n620\t# attn_mask = attention_mask.bool().view(B, 1, 1, L)\n621\t\n622\t# # Output: (B, Heads, K, HeadDim)\n623\t# attn_out = F.scaled_dot_product_attention(q, k, v, attn_mask=attn_mask)\n624\t \n625\t# # --- Reassemble ---\n626\t# # (B, Heads, K, HeadDim) -> (B, K, Heads, HeadDim) -> (B, K, D)\n627\t# attn_out = attn_out.transpose(1, 2).contiguous().view(B, K, D)\n628\t \n629\t# # Final Projection\n630\t# latents = self.o_proj(attn_out)\n631\t \n632\t# return latents\n633\t\n634\tclass PPIModelColBERT(nn.Module):\n635\t \"\"\"\n636\t Late Interaction (ColBERT-style) PPI Model.\n637\t Uses LatentPooling to output (B, K, D) embeddings instead of (B, L, D).\n638\t Uses shared PLM backbone like ContrastivePPIModel.\n639\t \"\"\"\n640\t def __init__(\n641\t self, \n642\t config: ModelConfig, \n643\t enable_contact_prediction: bool = True, \n644\t swap_self_negative: bool = False,\n645\t contact_random_neg_ratio: float = 1.0,\n646\t contact_self_neg_ratio: float = 1.0,\n647\t contact_hard_neg_ratio: float = 0.0,\n648\t valid_keys_tensor: Optional[torch.Tensor] = None,\n649\t multiplier: int = 1_000_000,\n650\t ):\n651\t pass\n652\t # super().__init__()\n653\t # self.config = config\n654\t # self.enable_contact_prediction = enable_contact_prediction\n655\t # self.swap_self_negative = swap_self_negative\n656\t # self.contact_random_neg_ratio = contact_random_neg_ratio\n657\t # self.contact_self_neg_ratio = contact_self_neg_ratio\n658\t\n659\t # # Shared PLM backbone\n660\t # self.plm = load_plm_model(config.plm_model_name)\n661\t \n662\t # # Freeze logic\n663\t # if config.num_frozen_layers > 0:\n664\t # freeze_plm_layers(self.plm, config.num_frozen_layers)\n665\t\n666\t # if hasattr(self.plm.config, 'hidden_size'):\n667\t # hidden_dim = self.plm.config.hidden_size\n668\t # else:\n669\t # hidden_dim = self.plm.config.dim\n670\t \n671\t # # ESM tokenizer adds cls/eos tokens.\n672\t # self._adds_special_tokens = \"esm\" in config.plm_model_name.lower()\n673\t \n674\t # # Latent Pooling (Perceiver IO style)\n675\t # self.latent_pool_q = LatentPooling(hidden_dim, num_latents=config.num_latents, output_dim=hidden_dim)\n676\t # self.latent_pool_k = LatentPooling(hidden_dim, num_latents=config.num_latents, output_dim=hidden_dim)\n677\t \n678\t # # Asymmetric projection Heads (MLP for better Q/K differentiation)\n679\t # self.head_q = MLPHead(hidden_dim, config.clip_embed_dim)\n680\t # self.head_k = MLPHead(hidden_dim, config.clip_embed_dim)\n681\t \n682\t # # Logit scale for ColBERT MaxSim, use 1 by default\n683\t # self.logit_scale = nn.Parameter(torch.ones([]))\n684\t\n685\t # # Contact Head (same as ContrastivePPIModel)\n686\t # self.contact_head = ContactHead(hidden_dim, config.contact_embed_dim, num_heads=8)\n687\t\n688\t # def encode_protein(\n689\t # self,\n690\t # input_ids: torch.Tensor,\n691\t # attention_mask: Optional[torch.Tensor] = None,\n692\t # ):\n693\t # \"\"\"\n694\t # Runs the shared plm model on tokenized protein sequences.\n695\t \n696\t # Args:\n697\t # input_ids (torch.Tensor): (B, L) - Token IDs from the tokenizer\n698\t # attention_mask (Optional[torch.Tensor]): (B, L) - Attention mask (1 for real tokens, 0 for padding)\n699\t \n700\t # Returns:\n701\t # residue_embeds (torch.Tensor): (B, L, D_plm) - Embeddings for each residue\n702\t # \"\"\"\n703\t # outputs = self.plm(input_ids=input_ids, attention_mask=attention_mask)\n704\t # residue_embeds = outputs.last_hidden_state\n705\t \n706\t # # Return residue embeddings directly (B, L, D)\n707\t # return residue_embeds\n708\t\n709\t # def predict_contacts(self, embed1: torch.Tensor, embed2: torch.Tensor, mask1: torch.Tensor, mask2: torch.Tensor):\n710\t # \"\"\"\n711\t # Predicts a contact map using ContactHead (same as ContrastivePPIModel).\n712\t \n713\t # Args:\n714\t # embed1 (torch.Tensor): (B, L1, D_plm) - Residue embeddings for sequence 1 (Query)\n715\t # embed2 (torch.Tensor): (B, L2, D_plm) - Residue embeddings for sequence 2 (Key)\n716\t # mask1 (torch.Tensor): (B, L1) - Attention mask for seq1\n717\t # mask2 (torch.Tensor): (B, L2) - Attention mask for seq2\n718\t \n719\t # Returns:\n720\t # contact_logits (torch.Tensor): (B, L1, L2) - Contact prediction logits\n721\t # contact_valid_mask (torch.Tensor): (B, L1, L2) - Valid mask for contact prediction\n722\t # \"\"\"\n723\t # return self.contact_head(embed1, embed2, mask1, mask2)\n724\t\n725\t # def encode_pair(\n726\t # self,\n727\t # seq1_input_ids: torch.Tensor,\n728\t # seq2_input_ids: torch.Tensor,\n729\t # seq1_attention_mask: Optional[torch.Tensor] = None,\n730\t # seq2_attention_mask: Optional[torch.Tensor] = None,\n731\t # ):\n732\t # \"\"\"\n733\t # Encodes a pair of protein sequences through PLM in a single batched call.\n734\t # Returns PLM-level representations (residue embeddings only for ColBERT).\n735\t \n736\t # Returns:\n737\t # residue_embed1 (torch.Tensor): (B, L1, D_plm) - Residue embeddings for sequence 1\n738\t # residue_embed2 (torch.Tensor): (B, L2, D_plm) - Residue embeddings for sequence 2\n739\t # \"\"\"\n740\t # B = seq1_input_ids.shape[0]\n741\t \n742\t # # Batch both sequences together for a single PLM call\n743\t # batched_input_ids = torch.cat([seq1_input_ids, seq2_input_ids], dim=0)\n744\t # if seq1_attention_mask is not None and seq2_attention_mask is not None:\n745\t # batched_attention_mask = torch.cat([seq1_attention_mask, seq2_attention_mask], dim=0)\n746\t # else:\n747\t # batched_attention_mask = None\n748\t \n749\t # # Single PLM call: (2*B, L, D_plm)\n750\t # batched_residue_embeds = self.encode_protein(batched_input_ids, batched_attention_mask)\n751\t \n752\t # # Split back into seq1 and seq2\n753\t # residue_embed1 = batched_residue_embeds[:B] # (B, L1, D_plm)\n754\t # residue_embed2 = batched_residue_embeds[B:] # (B, L2, D_plm)\n755\t \n756\t # return residue_embed1, residue_embed2\n757\t \n758\t # def embed_pair(\n759\t # self,\n760\t # seq1_input_ids: torch.Tensor,\n761\t # seq2_input_ids: torch.Tensor,\n762\t # seq1_attention_mask: Optional[torch.Tensor] = None,\n763\t # seq2_attention_mask: Optional[torch.Tensor] = None,\n764\t # ):\n765\t # \"\"\"\n766\t # Encodes a pair of proteins and returns final normalized CLIP embeddings.\n767\t # Optimized for GradCache (no contact prediction, single PLM call).\n768\t \n769\t # Returns:\n770\t # clip_embed1 (torch.Tensor): (B, K, D_clip) - Normalized CLIP latents for seq1 (Q projection)\n771\t # clip_embed2 (torch.Tensor): (B, K, D_clip) - Normalized CLIP latents for seq2 (K projection)\n772\t # \"\"\"\n773\t # res1, res2 = self.encode_pair(\n774\t # seq1_input_ids, seq2_input_ids, seq1_attention_mask, seq2_attention_mask\n775\t # )\n776\t \n777\t # # Standard projections: seq1 -> query latents, seq2 -> key latents\n778\t # latents1 = self.latent_pool_q(res1, seq1_attention_mask)\n779\t # latents2 = self.latent_pool_k(res2, seq2_attention_mask)\n780\t # clip_embed1 = F.normalize(self.head_q(latents1), p=2, dim=-1)\n781\t # clip_embed2 = F.normalize(self.head_k(latents2), p=2, dim=-1)\n782\t \n783\t # k1_identity = None\n784\t # q2_identity = None\n785\t # if self.swap_self_negative and self.training:\n786\t # latents1_k = self.latent_pool_k(res1, seq1_attention_mask) # K latents for seq1\n787\t # latents2_q = self.latent_pool_q(res2, seq2_attention_mask) # Q latents for seq2\n788\t # k1_identity = F.normalize(self.head_k(latents1_k), p=2, dim=-1) # K(A)\n789\t # q2_identity = F.normalize(self.head_q(latents2_q), p=2, dim=-1) # Q(B)\n790\t \n791\t # return clip_embed1, clip_embed2, k1_identity, q2_identity\n792\t\n793\t # def forward(self, seq1_input_ids, seq2_input_ids, seq1_attention_mask, seq2_attention_mask):\n794\t # \"\"\"\n795\t # Full Forward for Gradient Step.\n796\t # \"\"\"\n797\t # # Encode both proteins in a single PLM call\n798\t # res1, res2 = self.encode_pair(\n799\t # seq1_input_ids, seq2_input_ids, seq1_attention_mask, seq2_attention_mask\n800\t # )\n801\t\n802\t # # Standard projections: seq1 -> query latents, seq2 -> key latents\n803\t # latents1 = self.latent_pool_q(res1, seq1_attention_mask) # (B, K, D_plm)\n804\t # latents2 = self.latent_pool_k(res2, seq2_attention_mask) # (B, K, D_plm)\n805\t \n806\t # # Project & Normalize (ColBERT style)\n807\t # clip1 = F.normalize(self.head_q(latents1), p=2, dim=-1) # (B, K, D_clip)\n808\t # clip2 = F.normalize(self.head_k(latents2), p=2, dim=-1) # (B, K, D_clip)\n809\t\n810\t # # Compute identity projections for identity loss (always compute, use conditionally)\n811\t # latents1_k = self.latent_pool_k(res1, seq1_attention_mask) # K latents for seq1\n812\t # latents2_q = self.latent_pool_q(res2, seq2_attention_mask) # Q latents for seq2\n813\t # k1_identity = F.normalize(self.head_k(latents1_k), p=2, dim=-1) # K(A)\n814\t # q2_identity = F.normalize(self.head_q(latents2_q), p=2, dim=-1) # Q(B)\n815\t\n816\t # # Contacts\n817\t # contact_logits = None\n818\t # contact_valid_mask = None\n819\t # if self.enable_contact_prediction:\n820\t # # Slice out [CLS] and [SEP] tokens if the tokenizer added them\n821\t # if self._adds_special_tokens:\n822\t # res1_c = res1[:, 1:-1]\n823\t # res2_c = res2[:, 1:-1]\n824\t # mask1_c = seq1_attention_mask[:, 1:-1] if seq1_attention_mask is not None else None\n825\t # mask2_c = seq2_attention_mask[:, 1:-1] if seq2_attention_mask is not None else None\n826\t # else:\n827\t # res1_c = res1\n828\t # res2_c = res2\n829\t # mask1_c = seq1_attention_mask\n830\t # mask2_c = seq2_attention_mask\n831\t\n832\t # # Add negative pairs for contact prediction using configurable ratios\n833\t # res1_all, res2_all, mask1_all, mask2_all = get_contact_negatives(\n834\t # res1_c, res2_c, mask1_c, mask2_c,\n835\t # random_neg_ratio=self.contact_random_neg_ratio,\n836\t # self_neg_ratio=self.contact_self_neg_ratio,\n837\t # training=self.training\n838\t # )\n839\t # contact_logits, contact_valid_mask = self.predict_contacts(\n840\t # res1_all, res2_all, mask1_all, mask2_all\n841\t # )\n842\t \n843\t # return clip1, clip2, contact_logits, self.logit_scale.exp(), contact_valid_mask, k1_identity, q2_identity"}]},"uuid":"15a0a374-2e70-45e1-a03b-92457a5c4680","timestamp":"2026-06-01T15:28:21.769Z","sourceToolAssistantUUID":"b88a253a-19a3-42d0-932f-4e91525e7939","userType":"external","entrypoint":"claude-vscode","cwd":"/home/ubuntu/ppi_project","sessionId":"5a87a0cd-6802-4ac5-a708-29b2bf7ea563","version":"2.1.159","gitBranch":"afdb"} +{"parentUuid":"a73d39f6-30fb-4a2d-9cfb-1f2bc30d10e4","isSidechain":true,"promptId":"ae94540a-045a-4847-b16d-c2358b5c9db6","agentId":"a07e7799acd8be342","type":"user","message":{"role":"user","content":[{"tool_use_id":"toolu_01T9fSVBDYvoikLQDCLeoNwd","type":"tool_result","content":"1\t\"\"\"\n2\taccelerate launch --config_file configs_accelerate/cpu.yaml train.py configs_train/train_cpu.yaml\n3\t\"\"\"\n4\timport argparse\n5\timport os\n6\timport sys\n7\timport yaml\n8\timport logging\n9\timport torch\n10\timport numpy as np\n11\timport datasets\n12\timport transformers\n13\tfrom torch.utils.data import DataLoader\n14\tfrom accelerate import Accelerator, DistributedDataParallelKwargs\n15\tfrom accelerate.logging import get_logger\n16\tfrom accelerate.utils import set_seed\n17\tfrom datasets import load_dataset, concatenate_datasets\n18\tfrom safetensors.torch import load_file\n19\tfrom transformers import AutoTokenizer, get_cosine_with_min_lr_schedule_with_warmup, HfArgumentParser\n20\tfrom tqdm.auto import tqdm\n21\t\n22\tfrom config import ModelConfig, TrainingConfig\n23\tfrom model import ContrastivePPIModel, PPIModelColBERT\n24\tfrom loss import ClipLoss, ContactLoss, GradCacheClipLoss, PPIGradientCache, ColBERTLoss, ColBERTGradCacheLoss\n25\tfrom dataset import DataCollatorForPPI, load_retrieval_dataset, build_interaction_keys, BenchmarkCollator, ClusterSampler\n26\tfrom eval import evaluate, evaluate_retrieval, evaluate_benchmark\n27\t\n28\tlogger = get_logger(__name__)\n29\t\n30\tdef add_prefix_to_dict(d: dict, prefix: str) -> dict:\n31\t return {f\"{prefix}/{k}\": v for k, v in d.items()}\n32\t\n33\tdef save_config(config_dict, output_dir):\n34\t config_path = os.path.join(output_dir, \"config.yaml\")\n35\t with open(config_path, 'w') as f:\n36\t yaml.safe_dump(config_dict, f, default_flow_style=False)\n37\t\n38\tdef get_grouped_params(model, weight_decay):\n39\t # No decay, important for logit_scale\n40\t no_decay = [\"bias\", \"LayerNorm.weight\", \"logit_scale\", \"alpha\", \"beta\"]\n41\t \n42\t optimizer_grouped_parameters = [\n43\t {\n44\t \"params\": [\n45\t p for n, p in model.named_parameters() \n46\t if not any(nd in n for nd in no_decay) and p.requires_grad\n47\t ],\n48\t \"weight_decay\": weight_decay,\n49\t },\n50\t {\n51\t \"params\": [\n52\t p for n, p in model.named_parameters() \n53\t if any(nd in n for nd in no_decay) and p.requires_grad\n54\t ],\n55\t \"weight_decay\": 0.0,\n56\t },\n57\t ]\n58\t return optimizer_grouped_parameters\n59\t\n60\t\n61\tdef main(config_file, run_name=None):\n62\t # 1. Parse Configs\n63\t parser = HfArgumentParser((ModelConfig, TrainingConfig))\n64\t model_config, training_config = parser.parse_yaml_file(config_file)\n65\t\n66\t if isinstance(training_config.learning_rate, str):\n67\t training_config.learning_rate = float(training_config.learning_rate)\n68\t if run_name:\n69\t training_config.run_name = run_name\n70\t training_config.output_dir = os.path.join(training_config.output_dir, run_name)\n71\t\n72\t # 2. Accelerator Setup\n73\t ddp_kwargs = DistributedDataParallelKwargs(find_unused_parameters=not training_config.enable_contact_prediction)\n74\t accelerator = Accelerator(\n75\t log_with=\"wandb\" if training_config.use_wandb else None,\n76\t kwargs_handlers=[ddp_kwargs]\n77\t )\n78\t set_seed(training_config.seed)\n79\t os.makedirs(training_config.output_dir, exist_ok=True)\n80\t \n81\t # Configure logging to show output in terminal\n82\t logging.basicConfig(\n83\t format=\"%(asctime)s - %(levelname)s - %(name)s - %(message)s\",\n84\t datefmt=\"%m/%d/%Y %H:%M:%S\",\n85\t level=logging.INFO,\n86\t )\n87\t \n88\t # Add file handler for logging to file\n89\t if accelerator.is_main_process:\n90\t log_file = os.path.join(training_config.output_dir, \"training.log\")\n91\t file_handler = logging.FileHandler(log_file)\n92\t logger.logger.addHandler(file_handler)\n93\t \n94\t # Set verbosity for datasets and transformers\n95\t if accelerator.is_local_main_process:\n96\t transformers.utils.logging.set_verbosity_info()\n97\t else:\n98\t datasets.utils.logging.set_verbosity_error()\n99\t transformers.utils.logging.set_verbosity_error()\n100\t \n101\t # Log command and configuration\n102\t string_of_command = f\"{' '.join(sys.argv)}\"\n103\t logger.info(\"command: \" + string_of_command)\n104\t logger.info(accelerator.state, main_process_only=True)\n105\t logger.info(f\"Training/evaluation parameters {training_config}\")\n106\t logger.info(f\"Model parameters {model_config}\")\n107\t \n108\t if training_config.use_wandb:\n109\t init_kwargs = {}\n110\t if training_config.run_name:\n111\t init_kwargs[\"wandb\"] = {\"name\": training_config.run_name}\n112\t accelerator.init_trackers(\n113\t training_config.project_name,\n114\t config={**vars(model_config), **vars(training_config)},\n115\t init_kwargs=init_kwargs\n116\t )\n117\t\n118\t # 3. Data Loading\n119\t if accelerator.is_main_process:\n120\t logger.info(f\"Loading dataset from {training_config.dataset_path}\")\n121\t try:\n122\t hf_dataset = load_dataset(training_config.dataset_path)\n123\t except Exception as e:\n124\t logger.error(e)\n125\t return\n126\t\n127\t tokenizer = AutoTokenizer.from_pretrained(model_config.plm_model_name, trust_remote_code=True)\n128\t if training_config.train_split == \"all\":\n129\t print(f\"Concatenating all splits: {hf_dataset.keys()}\")\n130\t train_dataset = concatenate_datasets(hf_dataset.values())\n131\t else:\n132\t train_dataset = hf_dataset[training_config.train_split]\n133\t \n134\t val_dataset = hf_dataset[training_config.val_split].shuffle(seed=training_config.seed)\n135\t\n136\t # Filter homomers if enabled (only for main dataset, not DDI)\n137\t if training_config.filter_homomers:\n138\t with accelerator.main_process_first():\n139\t initial_size = len(train_dataset)\n140\t train_dataset = train_dataset.filter(\n141\t lambda ex: not ((ex.get('uniprot1_id') == ex.get('uniprot2_id') and ex.get('uniprot1_id') is not None) or \n142\t (ex.get('seq1_cluster') == ex.get('seq2_cluster') and ex.get('seq1_cluster') is not None and ex.get('seq1_cluster') != -1)),\n143\t num_proc=training_config.dataloader_num_workers\n144\t )\n145\t if accelerator.is_main_process:\n146\t logger.info(f\"Filtered {initial_size - len(train_dataset)} homomers. Remaining: {len(train_dataset)} examples\")\n147\t\n148\t # Truncate validation dataset if max_eval_examples is set\n149\t if training_config.max_eval_examples > 0:\n150\t max_examples = min(training_config.max_eval_examples, len(val_dataset))\n151\t val_dataset = val_dataset.select(range(max_examples))\n152\t\n153\t # Extract cluster columns\n154\t c1_train = np.array(train_dataset['seq1_cluster'], dtype=np.int64)\n155\t c2_train = np.array(train_dataset['seq2_cluster'], dtype=np.int64)\n156\t c1_val = np.array(val_dataset['seq1_cluster'], dtype=np.int64)\n157\t c2_val = np.array(val_dataset['seq2_cluster'], dtype=np.int64)\n158\t\n159\t # Mark as positive interactions\n160\t train_dataset = train_dataset.add_column(\"is_positive\", [True] * len(train_dataset))\n161\t\n162\t # Track concatenated index ranges for each extra pool (used by ClusterSampler)\n163\t ddi_pool_indices = np.array([], dtype=np.int64)\n164\t afdb_pool_indices = np.array([], dtype=np.int64)\n165\t\n166\t # Load and concatenate DDI dataset if specified\n167\t if training_config.ddi_dataset is not None:\n168\t if accelerator.is_main_process:\n169\t logger.info(f\"Loading DDI dataset from {training_config.ddi_dataset}\")\n170\t ddi_hf_dataset = load_dataset(training_config.ddi_dataset)\n171\t ddi_train = ddi_hf_dataset['train']\n172\t c1_ddi = np.full(len(ddi_train), -1, dtype=np.int64)\n173\t c2_ddi = np.full(len(ddi_train), -1, dtype=np.int64)\n174\t # DDI interactions are positives\n175\t ddi_train = ddi_train.add_column(\"is_positive\", [True] * len(ddi_train))\n176\t \n177\t if training_config.ddi_only:\n178\t if accelerator.is_main_process:\n179\t logger.info(f\"Using DDI dataset only ({len(ddi_train)} examples)\")\n180\t train_dataset = ddi_train\n181\t c1_train = c1_ddi\n182\t c2_train = c2_ddi\n183\t ddi_pool_indices = np.arange(len(ddi_train), dtype=np.int64)\n184\t else:\n185\t if accelerator.is_main_process:\n186\t logger.info(f\"Concatenating DDI dataset ({len(ddi_train)} examples) with train dataset ({len(train_dataset)} examples)\")\n187\t ddi_start = len(train_dataset)\n188\t train_dataset = concatenate_datasets([train_dataset, ddi_train])\n189\t c1_train = np.concatenate([c1_train, c1_ddi])\n190\t c2_train = np.concatenate([c2_train, c2_ddi])\n191\t ddi_pool_indices = np.arange(ddi_start, ddi_start + len(ddi_train), dtype=np.int64)\n192\t\n193\t # Load and concatenate AlphaFold heterodimer dataset if specified\n194\t if training_config.afdb_dataset is not None:\n195\t if accelerator.is_main_process:\n196\t logger.info(f\"Loading AFDB dataset from {training_config.afdb_dataset}\")\n197\t afdb_hf_dataset = load_dataset(training_config.afdb_dataset)\n198\t afdb_train = afdb_hf_dataset['train']\n199\t c1_afdb = np.full(len(afdb_train), -1, dtype=np.int64)\n200\t c2_afdb = np.full(len(afdb_train), -1, dtype=np.int64)\n201\t afdb_train = afdb_train.add_column(\"is_positive\", [True] * len(afdb_train))\n202\t # Ensure schema compatibility before concatenation:\n203\t # 1. Add missing cluster columns (AFDB has no cluster IDs).\n204\t for col in ['seq1_cluster', 'seq2_cluster']:\n205\t if col not in afdb_train.column_names and col in train_dataset.column_names:\n206\t afdb_train = afdb_train.add_column(col, [-1] * len(afdb_train))\n207\t # 2. Cast shared columns to match train_dataset types (e.g. large_string -> string).\n208\t cast_features = datasets.Features({\n209\t col: (train_dataset.features[col] if col in train_dataset.features else feat)\n210\t for col, feat in afdb_train.features.items()\n211\t })\n212\t afdb_train = afdb_train.cast(cast_features)\n213\t if accelerator.is_main_process:\n214\t logger.info(f\"Concatenating AFDB dataset ({len(afdb_train)} examples) with train dataset ({len(train_dataset)} examples)\")\n215\t afdb_start = len(train_dataset)\n216\t train_dataset = concatenate_datasets([train_dataset, afdb_train])\n217\t c1_train = np.concatenate([c1_train, c1_afdb])\n218\t c2_train = np.concatenate([c2_train, c2_afdb])\n219\t afdb_pool_indices = np.arange(afdb_start, afdb_start + len(afdb_train), dtype=np.int64)\n220\t\n221\t # Load and concatenate Negative dataset if specified\n222\t if training_config.neg_dataset is not None:\n223\t if accelerator.is_main_process:\n224\t logger.info(f\"Loading negative dataset from {training_config.neg_dataset}\")\n225\t neg_hf_dataset = load_dataset(training_config.neg_dataset)\n226\t neg_train = neg_hf_dataset['train']\n227\t # Negative interactions (no contacts) are False for CLIP positive\n228\t neg_train = neg_train.add_column(\"is_positive\", [False] * len(neg_train))\n229\t \n230\t if accelerator.is_main_process:\n231\t logger.info(f\"Concatenating negative dataset ({len(neg_train)} examples) with train dataset ({len(train_dataset)} examples)\")\n232\t train_dataset = concatenate_datasets([train_dataset, neg_train])\n233\t \n234\t # ESM tokenizer adds cls/eos tokens.\n235\t adds_special_tokens = \"esm\" in model_config.plm_model_name.lower()\n236\t train_collator = DataCollatorForPPI(\n237\t tokenizer, \n238\t model_config.max_len, \n239\t adds_special_tokens=adds_special_tokens\n240\t )\n241\t val_collator = DataCollatorForPPI(tokenizer, model_config.max_len, swap_augment=False, adds_special_tokens=adds_special_tokens)\n242\t \n243\t # Build cluster sampler if enabled\n244\t train_sampler = None\n245\t if training_config.use_cluster_weighted_sampling:\n246\t if accelerator.is_main_process:\n247\t logger.info(\"Building cluster sampler...\")\n248\t\n249\t ppi_weight = 1.0 - training_config.ddi_sampling_weight - training_config.afdb_sampling_weight\n250\t extra_pools = []\n251\t sampling_probs = [ppi_weight]\n252\t if training_config.ddi_dataset is not None:\n253\t extra_pools.append(ddi_pool_indices)\n254\t sampling_probs.append(training_config.ddi_sampling_weight)\n255\t if training_config.afdb_dataset is not None:\n256\t extra_pools.append(afdb_pool_indices)\n257\t sampling_probs.append(training_config.afdb_sampling_weight)\n258\t\n259\t train_sampler = ClusterSampler(\n260\t c1_train, c2_train,\n261\t extra_pools=extra_pools,\n262\t sampling_probs=sampling_probs,\n263\t seed=training_config.seed\n264\t )\n265\t if accelerator.is_main_process:\n266\t pool_names = ([\"DDI\"] + [\"AFDB\"])[:len(extra_pools)]\n267\t pool_info = \", \".join(\n268\t f\"{name}={len(pool):,} examples ({prob:.1%})\"\n269\t for name, pool, prob in zip(pool_names, extra_pools, sampling_probs[1:])\n270\t )\n271\t logger.info(\n272\t f\"Cluster sampling: {len(train_sampler.clusters)} PPI clusters ({ppi_weight:.1%}), {pool_info}\"\n273\t )\n274\t \n275\t # IMPORTANT: batch_size here is the LARGE effective batch size (e.g., 256)\n276\t train_loader = DataLoader(\n277\t train_dataset,\n278\t batch_size=training_config.batch_size,\n279\t collate_fn=train_collator,\n280\t sampler=train_sampler,\n281\t shuffle=(train_sampler is None), # Only shuffle if no sampler\n282\t num_workers=training_config.dataloader_num_workers,\n283\t pin_memory=True\n284\t )\n285\t \n286\t # Eval uses smaller batches\n287\t val_loader = DataLoader(\n288\t val_dataset,\n289\t batch_size=training_config.eval_batch_size,\n290\t collate_fn=val_collator,\n291\t shuffle=False,\n292\t num_workers=training_config.dataloader_num_workers,\n293\t pin_memory=True\n294\t )\n295\t \n296\t # Build Global Interaction Map (for masking false negatives)\n297\t if accelerator.is_main_process:\n298\t logger.info(\"Building global interaction map from train and validation datasets...\")\n299\t valid_keys, multiplier = build_interaction_keys(c1_train, c2_train, c1_val, c2_val)\n300\t\n301\t # 4. Model & Loss & Optimizer\n302\t if model_config.model_type == \"late_retrieval\":\n303\t model = PPIModelColBERT(\n304\t model_config, \n305\t enable_contact_prediction=training_config.enable_contact_prediction,\n306\t swap_self_negative=training_config.swap_self_negative,\n307\t contact_random_neg_ratio=training_config.contact_random_neg_ratio,\n308\t contact_self_neg_ratio=training_config.contact_self_neg_ratio,\n309\t contact_hard_neg_ratio=training_config.contact_hard_neg_ratio,\n310\t valid_keys_tensor=valid_keys,\n311\t multiplier=multiplier,\n312\t )\n313\t # Use MaxSim Loss\n314\t train_loss_fn = ColBERTGradCacheLoss(\n315\t rank=accelerator.process_index, \n316\t world_size=accelerator.num_processes,\n317\t valid_keys_tensor=valid_keys,\n318\t multiplier=multiplier\n319\t )\n320\t else:\n321\t model = ContrastivePPIModel(\n322\t model_config, \n323\t enable_contact_prediction=training_config.enable_contact_prediction,\n324\t swap_self_negative=training_config.swap_self_negative,\n325\t enable_clip=training_config.enable_clip,\n326\t contact_random_neg_ratio=training_config.contact_random_neg_ratio,\n327\t contact_self_neg_ratio=training_config.contact_self_neg_ratio,\n328\t contact_hard_neg_ratio=training_config.contact_hard_neg_ratio,\n329\t valid_keys_tensor=valid_keys,\n330\t multiplier=multiplier,\n331\t )\n332\t train_loss_fn = GradCacheClipLoss(\n333\t rank=accelerator.process_index, \n334\t world_size=accelerator.num_processes,\n335\t valid_keys_tensor=valid_keys,\n336\t multiplier=multiplier\n337\t )\n338\t\n339\t optimizer = torch.optim.AdamW(\n340\t get_grouped_params(model, training_config.weight_decay),\n341\t lr=training_config.learning_rate\n342\t )\n343\t\n344\t num_steps = len(train_loader) * training_config.max_epochs\n345\t lr_scheduler = get_cosine_with_min_lr_schedule_with_warmup(optimizer, int(num_steps * training_config.warmup_fraction), num_steps, min_lr_rate=0.1)\n346\t\n347\t # 5. Load retrieval evaluation dataset (if enabled)\n348\t retrieval_data = None\n349\t if training_config.retrieval_eval_path:\n350\t retrieval_data = load_retrieval_dataset(\n351\t training_config.retrieval_eval_path,\n352\t tokenizer,\n353\t model_config.max_len\n354\t )\n355\t \n356\t # Load benchmark evaluation dataset (if enabled) - load on all processes for DDP\n357\t benchmark_loader = None\n358\t if training_config.benchmark_eval_path:\n359\t with accelerator.main_process_first():\n360\t benchmark_hf_dataset = load_dataset(training_config.benchmark_eval_path)\n361\t # Get the split name (should have only one split)\n362\t split_names = list(benchmark_hf_dataset.keys())\n363\t assert len(split_names) == 1, f\"Expected exactly one split in benchmark dataset, but found {len(split_names)}: {split_names}\"\n364\t benchmark_dataset = benchmark_hf_dataset[split_names[0]]\n365\t # Add index column to track original dataset order\n366\t benchmark_dataset = benchmark_dataset.add_column(\"index\", range(len(benchmark_dataset)))\n367\t if accelerator.is_main_process:\n368\t logger.info(f\"Loaded benchmark dataset with {len(benchmark_dataset)} examples\")\n369\t benchmark_collator = BenchmarkCollator(tokenizer, 1024, adds_special_tokens)\n370\t benchmark_loader = DataLoader(\n371\t benchmark_dataset,\n372\t batch_size=training_config.eval_batch_size,\n373\t collate_fn=benchmark_collator,\n374\t shuffle=False,\n375\t num_workers=training_config.dataloader_num_workers,\n376\t pin_memory=True\n377\t )\n378\t\n379\t # 6. Load model weights only (BEFORE prepare) if specified\n380\t if training_config.load_from_checkpoint:\n381\t model_path = os.path.join(training_config.load_from_checkpoint, \"model.safetensors\")\n382\t if os.path.exists(model_path):\n383\t state_dict = load_file(model_path)\n384\t model.load_state_dict(state_dict, strict=False)\n385\t if accelerator.is_main_process:\n386\t logger.info(f\"Loaded model weights from: {model_path}\")\n387\t else:\n388\t raise FileNotFoundError(f\"Model weights not found at {model_path}\")\n389\t \n390\t # Enable gradient checkpointing if specified in config\n391\t if model_config.gradient_checkpointing:\n392\t model.gradient_checkpointing_enable()\n393\t if accelerator.is_main_process:\n394\t logger.info(\"Gradient checkpointing enabled\")\n395\t \n396\t # Prepare\n397\t if benchmark_loader is not None:\n398\t model, optimizer, train_loader, val_loader, lr_scheduler, benchmark_loader = accelerator.prepare(\n399\t model, optimizer, train_loader, val_loader, lr_scheduler, benchmark_loader\n400\t )\n401\t else:\n402\t model, optimizer, train_loader, val_loader, lr_scheduler = accelerator.prepare(\n403\t model, optimizer, train_loader, val_loader, lr_scheduler\n404\t )\n405\t\n406\t # Load full training state (model + optimizer + scheduler) if specified\n407\t if training_config.resume_from_checkpoint:\n408\t accelerator.load_state(training_config.resume_from_checkpoint)\n409\t if accelerator.is_main_process:\n410\t logger.info(f\"Resumed full training state from checkpoint: {training_config.resume_from_checkpoint}\")\n411\t\n412\t # 7. Create config dict for saving (once before training)\n413\t config_dict = {**{k: v for k, v in vars(model_config).items()}, \n414\t **{k: v for k, v in vars(training_config).items()}}\n415\t \n416\t # 8. Initialize Loss Components\n417\t if model_config.model_type == \"late_retrieval\":\n418\t eval_loss_fn = ColBERTLoss(\n419\t local_loss=True, \n420\t gather_with_grad=False, \n421\t rank=accelerator.process_index, \n422\t world_size=accelerator.num_processes,\n423\t valid_keys_tensor=valid_keys,\n424\t multiplier=multiplier\n425\t )\n426\t else:\n427\t eval_loss_fn = ClipLoss(\n428\t local_loss=True, \n429\t gather_with_grad=False, \n430\t rank=accelerator.process_index, \n431\t world_size=accelerator.num_processes,\n432\t valid_keys_tensor=valid_keys,\n433\t multiplier=multiplier\n434\t )\n435\t contact_loss = ContactLoss(use_focal_loss=training_config.use_focal_loss)\n436\t \n437\t # Conditionally initialize GradCache or standard training loss\n438\t use_grad_cache = training_config.grad_cache_chunk_size is not None\n439\t if use_grad_cache:\n440\t grad_cache = PPIGradientCache(\n441\t model=model,\n442\t chunk_size=training_config.grad_cache_chunk_size,\n443\t clip_loss_fn=train_loss_fn,\n444\t contact_loss_fn=contact_loss,\n445\t contact_weight=training_config.contact_loss_weight if training_config.enable_contact_prediction else 0.0,\n446\t accelerator=accelerator,\n447\t enable_clip=training_config.enable_clip,\n448\t clip_weight=training_config.clip_loss_weight if training_config.enable_clip else 0.0\n449\t )\n450\t else:\n451\t # Standard training loss with gradient gathering for distributed training\n452\t if model_config.model_type == \"late_retrieval\":\n453\t train_loss_fn = ColBERTLoss(\n454\t rank=accelerator.process_index, \n455\t world_size=accelerator.num_processes,\n456\t valid_keys_tensor=valid_keys,\n457\t multiplier=multiplier\n458\t )\n459\t else:\n460\t train_loss_fn = ClipLoss(\n461\t local_loss=False, \n462\t gather_with_grad=True, \n463\t rank=accelerator.process_index, \n464\t world_size=accelerator.num_processes,\n465\t valid_keys_tensor=valid_keys,\n466\t multiplier=multiplier\n467\t )\n468\t grad_cache = None\n469\t\n470\t # 7. Training Loop\n471\t if use_grad_cache:\n472\t logger.info(\"***** Running Training with Custom GradCache *****\")\n473\t logger.info(f\" Effective Batch: {training_config.batch_size * accelerator.num_processes}\")\n474\t logger.info(f\" Physical Chunk: {training_config.grad_cache_chunk_size}\")\n475\t else:\n476\t logger.info(\"***** Running Training with Standard Loss *****\")\n477\t logger.info(f\" Batch Size: {training_config.batch_size * accelerator.num_processes}\")\n478\t\n479\t completed_steps = 0\n480\t \n481\t for epoch in range(training_config.max_epochs):\n482\t model.train()\n483\t # Only call set_epoch if the sampler supports it (for DDP reproducibility)\n484\t if hasattr(train_loader.sampler, 'set_epoch'):\n485\t train_loader.sampler.set_epoch(epoch)\n486\t progress_bar = tqdm(train_loader, disable=not accelerator.is_local_main_process, desc=f\"Epoch {epoch+1}\")\n487\t\n488\t for step, batch in enumerate(progress_bar):\n489\t if use_grad_cache:\n490\t # The .step() method handles Cache->Loss->Gradient workflow\n491\t # It handles backward() internally\n492\t clip_loss_val, contact_loss_val, logit_scale_val, top1_acc = grad_cache.step(batch)\n493\t else:\n494\t # Standard loss computation\n495\t seq1_ids = batch['seq1_tokens']['input_ids']\n496\t seq1_mask = batch['seq1_tokens']['attention_mask']\n497\t seq2_ids = batch['seq2_tokens']['input_ids']\n498\t seq2_mask = batch['seq2_tokens']['attention_mask']\n499\t contact_target = batch['contact_target']\n500\t seq1_clusters = batch.get('seq1_cluster', None)\n501\t seq2_clusters = batch.get('seq2_cluster', None)\n502\t \n503\t # Forward pass (support both ContrastivePPIModel and PPIModelColBERT outputs)\n504\t # Model handles cluster ID swapping internally when swap_self_negative is enabled\n505\t outputs = model(\n506\t seq1_ids, seq2_ids, seq1_mask, seq2_mask,\n507\t seq1_clusters, seq2_clusters\n508\t )\n509\t clip_embed1, clip_embed2, contact_logits, logit_scale, contact_valid_mask, seq1_clusters, seq2_clusters = outputs[:7]\n510\t \n511\t if training_config.enable_clip:\n512\t clip_loss, top1_acc = train_loss_fn(clip_embed1, clip_embed2, logit_scale, seq1_clusters=seq1_clusters, seq2_clusters=seq2_clusters, compute_top1_acc=True)\n513\t clip_loss_val = clip_loss.item()\n514\t else:\n515\t clip_loss = torch.tensor(0.0, device=accelerator.device)\n516\t clip_loss_val = 0.0\n517\t top1_acc = 0.0\n518\t\n519\t contact_loss_val = 0.0\n520\t if training_config.enable_contact_prediction and contact_logits is not None:\n521\t contact_loss_tensor = contact_loss(contact_logits, contact_target, contact_valid_mask)\n522\t contact_loss_val = contact_loss_tensor.item()\n523\t total_loss = training_config.clip_loss_weight * clip_loss + training_config.contact_loss_weight * contact_loss_tensor\n524\t else:\n525\t total_loss = training_config.clip_loss_weight * clip_loss\n526\t \n527\t logit_scale_val = logit_scale.item()\n528\t \n529\t # Backward pass\n530\t accelerator.backward(total_loss)\n531\t \n532\t # Gradient Clipping\n533\t grad_norm = None\n534\t if training_config.max_grad_norm is not None:\n535\t grad_norm = accelerator.clip_grad_norm_(model.parameters(), training_config.max_grad_norm)\n536\t if isinstance(grad_norm, torch.Tensor):\n537\t grad_norm = grad_norm.item()\n538\t\n539\t # Optimizer Step\n540\t optimizer.step()\n541\t lr_scheduler.step()\n542\t optimizer.zero_grad()\n543\t completed_steps += 1\n544\t \n545\t # Logging\n546\t if completed_steps % training_config.log_interval == 0:\n547\t log_dict = {\n548\t \"train/loss_clip\": clip_loss_val,\n549\t \"train/loss_contact\": contact_loss_val,\n550\t \"train/loss_total\": clip_loss_val + contact_loss_val,\n551\t \"train/lr\": lr_scheduler.get_last_lr()[0],\n552\t \"train/logit_scale\": logit_scale_val,\n553\t \"train/retrieval_top1_acc\": top1_acc\n554\t }\n555\t if grad_norm is not None:\n556\t log_dict[\"train/grad_norm\"] = grad_norm\n557\t accelerator.log(log_dict, step=completed_steps)\n558\t progress_bar.set_postfix({\n559\t \"clip\": clip_loss_val, \n560\t \"contact\": contact_loss_val,\n561\t \"R@1\": f\"{top1_acc:.3f}\"\n562\t })\n563\t \n564\t # Saving\n565\t if completed_steps % training_config.save_interval == 0 and accelerator.is_main_process:\n566\t checkpoint_dir = os.path.join(training_config.output_dir, f\"step_{completed_steps}\")\n567\t accelerator.save_state(checkpoint_dir)\n568\t save_config(config_dict, checkpoint_dir)\n569\t \n570\t # Evaluation\n571\t if completed_steps % training_config.eval_interval == 0:\n572\t # Re-instantiate standard CLIP loss for eval loop if needed or use the one prepared\n573\t evaluate(model, val_loader, accelerator, training_config, eval_loss_fn, contact_loss, completed_steps, adds_special_tokens)\n574\t \n575\t # Retrieval evaluation (if enabled and CLIP is enabled)\n576\t if training_config.enable_clip and retrieval_data and accelerator.is_main_process:\n577\t queries = retrieval_data['queries']\n578\t corpus = retrieval_data['corpus']\n579\t query_inputs = retrieval_data['query_inputs']\n580\t corpus_inputs = retrieval_data['corpus_inputs']\n581\t \n582\t query_ids = [q['binder_uniprot_id'] for q in queries]\n583\t corpus_ids = [c['uniprot_id'] for c in corpus]\n584\t query_complex_ids = [q.get('complex_id', '') for q in queries]\n585\t query_contacts = [q.get('contacts', []) for q in queries]\n586\t query_sequences = [q['sequence'] for q in queries] if query_contacts else None\n587\t corpus_sequences = [c['sequence'] for c in corpus] if query_contacts else None\n588\t query_all_binder_ids = [q.get('all_binder_uniprot_ids', [q['binder_uniprot_id']]) for q in queries]\n589\t query_uniprot_ids = [q.get('uniprot_id') for q in queries]\n590\t \n591\t retrieval_metrics, _ = evaluate_retrieval(\n592\t model=model,\n593\t batch_size=training_config.eval_batch_size,\n594\t query_inputs=query_inputs,\n595\t corpus_inputs=corpus_inputs,\n596\t query_ids=query_ids,\n597\t corpus_ids=corpus_ids,\n598\t query_complex_ids=query_complex_ids,\n599\t contact_metrics_max_examples=training_config.contact_metrics_max_examples,\n600\t device=accelerator.device,\n601\t query_contacts=query_contacts,\n602\t query_sequences=query_sequences,\n603\t corpus_sequences=corpus_sequences,\n604\t max_len=model_config.max_len,\n605\t model_type=model_config.model_type,\n606\t query_all_binder_ids=query_all_binder_ids,\n607\t query_uniprot_ids=query_uniprot_ids,\n608\t adds_special_tokens=adds_special_tokens\n609\t )\n610\t accelerator.log({f\"retrieval/{k}\": v for k, v in retrieval_metrics.items()}, step=completed_steps)\n611\t\n612\t # Benchmark evaluation (if enabled)\n613\t if benchmark_loader is not None:\n614\t benchmark_metrics = evaluate_benchmark(\n615\t model=model,\n616\t benchmark_loader=benchmark_loader,\n617\t accelerator=accelerator,\n618\t )\n619\t if benchmark_metrics:\n620\t accelerator.log({f\"benchmark/{k}\": v for k, v in benchmark_metrics.items()}, step=completed_steps)\n621\t \n622\t model.train()\n623\t\n624\t accelerator.wait_for_everyone()\n625\t if accelerator.is_main_process:\n626\t checkpoint_dir = os.path.join(training_config.output_dir, f\"epoch_{epoch+1}\")\n627\t accelerator.save_model(model, checkpoint_dir)\n628\t save_config(config_dict, checkpoint_dir)\n629\t\n630\t accelerator.end_training()\n631\t\n632\tif __name__ == \"__main__\": \n633\t parser = argparse.ArgumentParser()\n634\t parser.add_argument(\"config_file\", type=str, default=\"train.yaml\", nargs=\"?\")\n635\t parser.add_argument(\"--run_name\", type=str, default=None)\n636\t args = parser.parse_args() \n637\t main(args.config_file, run_name=args.run_name)"}]},"uuid":"76c424d4-d9f6-4f0b-941e-84d760517b7d","timestamp":"2026-06-01T15:28:21.818Z","sourceToolAssistantUUID":"a73d39f6-30fb-4a2d-9cfb-1f2bc30d10e4","userType":"external","entrypoint":"claude-vscode","cwd":"/home/ubuntu/ppi_project","sessionId":"5a87a0cd-6802-4ac5-a708-29b2bf7ea563","version":"2.1.159","gitBranch":"afdb"} +{"parentUuid":"76c424d4-d9f6-4f0b-941e-84d760517b7d","isSidechain":true,"agentId":"a07e7799acd8be342","message":{"model":"claude-haiku-4-5-20251001","id":"msg_01SFSXP6F8RaZ7D9jC5i34mF","type":"message","role":"assistant","content":[{"type":"text","text":"Now let me search for additional data-related files and check for any windowing/cropping logic:"}],"stop_reason":null,"stop_sequence":null,"stop_details":null,"usage":{"input_tokens":13825,"cache_creation_input_tokens":15130,"cache_read_input_tokens":26543,"cache_creation":{"ephemeral_5m_input_tokens":15130,"ephemeral_1h_input_tokens":0},"output_tokens":1,"service_tier":"standard","inference_geo":"not_available"},"diagnostics":null},"requestId":"req_011CbcmzSP4j7vMv6uJsT2EM","attributionAgent":"Explore","type":"assistant","uuid":"a5c789ce-e0f1-47b5-97b1-8e35de3cda29","timestamp":"2026-06-01T15:28:23.535Z","userType":"external","entrypoint":"claude-vscode","cwd":"/home/ubuntu/ppi_project","sessionId":"5a87a0cd-6802-4ac5-a708-29b2bf7ea563","version":"2.1.159","gitBranch":"afdb"} +{"parentUuid":"a5c789ce-e0f1-47b5-97b1-8e35de3cda29","isSidechain":true,"agentId":"a07e7799acd8be342","message":{"model":"claude-haiku-4-5-20251001","id":"msg_01SFSXP6F8RaZ7D9jC5i34mF","type":"message","role":"assistant","content":[{"type":"tool_use","id":"toolu_013TfiJZfnEiRq7RwiCSn7HX","name":"Bash","input":{"command":"find /home/ubuntu/ppi_project -type f -name \"*.py\" -path \"*/ppi_project/*\" ! -path \"*/ppi-env/*\" | grep -v __pycache__ | head -50"},"caller":{"type":"direct"}}],"stop_reason":null,"stop_sequence":null,"stop_details":null,"usage":{"input_tokens":13825,"cache_creation_input_tokens":15130,"cache_read_input_tokens":26543,"cache_creation":{"ephemeral_5m_input_tokens":15130,"ephemeral_1h_input_tokens":0},"output_tokens":1,"service_tier":"standard","inference_geo":"not_available"},"diagnostics":null},"requestId":"req_011CbcmzSP4j7vMv6uJsT2EM","attributionAgent":"Explore","type":"assistant","uuid":"b99a7cf4-4451-4ccd-81a2-467c234baa77","timestamp":"2026-06-01T15:28:24.382Z","userType":"external","entrypoint":"claude-vscode","cwd":"/home/ubuntu/ppi_project","sessionId":"5a87a0cd-6802-4ac5-a708-29b2bf7ea563","version":"2.1.159","gitBranch":"afdb"} +{"parentUuid":"b99a7cf4-4451-4ccd-81a2-467c234baa77","isSidechain":true,"promptId":"ae94540a-045a-4847-b16d-c2358b5c9db6","agentId":"a07e7799acd8be342","type":"user","message":{"role":"user","content":[{"tool_use_id":"toolu_013TfiJZfnEiRq7RwiCSn7HX","type":"tool_result","content":"/home/ubuntu/ppi_project/evaluate_rf2_benchmark.py\n/home/ubuntu/ppi_project/evaluate_rf2_benchmark_2model.py\n/home/ubuntu/ppi_project/eval.py\n/home/ubuntu/ppi_project/model.py\n/home/ubuntu/ppi_project/process_pdb_contacts.py\n/home/ubuntu/ppi_project/upload_ds_to_hf.py\n/home/ubuntu/ppi_project/eval_checkpoint.py\n/home/ubuntu/ppi_project/inference_genome_2stage.py\n/home/ubuntu/ppi_project/loss.py\n/home/ubuntu/ppi_project/run_colabfold.py\n/home/ubuntu/ppi_project/run_plm_interact_benchmark.py\n/home/ubuntu/ppi_project/inference_laub_virus_ecoli.py\n/home/ubuntu/ppi_project/find_novel_ppi.py\n/home/ubuntu/ppi_project/inference_genome.py\n/home/ubuntu/ppi_project/push_model_to_hf.py\n/home/ubuntu/ppi_project/layers.py\n/home/ubuntu/ppi_project/train.py\n/home/ubuntu/ppi_project/plot_boltz_vs_flashppi.py\n/home/ubuntu/ppi_project/inference_virushostdb.py\n/home/ubuntu/ppi_project/extract_queries.py\n/home/ubuntu/ppi_project/config.py\n/home/ubuntu/ppi_project/negative_control_virushostdb.py\n/home/ubuntu/ppi_project/dataset.py\n/home/ubuntu/ppi_project/plot_afmm_vs_flashppi_virhost.py\n/home/ubuntu/ppi_project/evaluate_pdb_benchmark.py\n/home/ubuntu/ppi_project/plot_afmm_vs_flashppi_ecoli.py\n/home/ubuntu/ppi_project/scripts/plot_stage1_retrieval_recall.py\n/home/ubuntu/ppi_project/scripts/create_ecoli_pdb_benchmark.py\n/home/ubuntu/ppi_project/scripts/compute_precision_thresholds.py\n/home/ubuntu/ppi_project/scripts/remove_leakage_ecoli_pdb_benchmark.py\n/home/ubuntu/ppi_project/scripts/plot_m_genitalium_iptm_predictions.py\n/home/ubuntu/ppi_project/scripts/plot_string_precision_recall_k.py\n/home/ubuntu/ppi_project/scripts/modal_alphafold.py\n/home/ubuntu/ppi_project/scripts/plot_virhost_negative_controls.py\n/home/ubuntu/ppi_project/scripts/plot_ecoli_pdb_benchmark_results.py\n/home/ubuntu/ppi_project/scripts/plot_example_contacts.py\n/home/ubuntu/ppi_project/scripts/evaluate_rf2_benchmark_dscript.py\n/home/ubuntu/ppi_project/scripts/run_msapairformer.py\n/home/ubuntu/ppi_project/scripts/cluster_dataset.py\n/home/ubuntu/ppi_project/scripts/plot_time.py\n/home/ubuntu/ppi_project/scripts/plot_string_other_contacts.py\n/home/ubuntu/ppi_project/scripts/plot_venn_string_predictions.py\n/home/ubuntu/ppi_project/scripts/plot_contact_metrics.py\n/home/ubuntu/ppi_project/scripts/build_test_set.py\n/home/ubuntu/ppi_project/hf_model/__init__.py\n/home/ubuntu/ppi_project/hf_model/configuration_flashppi.py\n/home/ubuntu/ppi_project/hf_model/modeling_flashppi.py\n/home/ubuntu/ppi_project/plm_models/modeling_esm_flash.py\n/home/ubuntu/ppi_project/plm_models/configuration_glm2.py\n/home/ubuntu/ppi_project/plm_models/modeling_glm2_flash.py","is_error":false}]},"uuid":"a2527e55-71f2-4a49-8059-89e9d301aab0","timestamp":"2026-06-01T15:28:24.444Z","sourceToolAssistantUUID":"b99a7cf4-4451-4ccd-81a2-467c234baa77","userType":"external","entrypoint":"claude-vscode","cwd":"/home/ubuntu/ppi_project","sessionId":"5a87a0cd-6802-4ac5-a708-29b2bf7ea563","version":"2.1.159","gitBranch":"afdb"} +{"parentUuid":"a2527e55-71f2-4a49-8059-89e9d301aab0","isSidechain":true,"agentId":"a07e7799acd8be342","message":{"model":"claude-haiku-4-5-20251001","id":"msg_01SFSXP6F8RaZ7D9jC5i34mF","type":"message","role":"assistant","content":[{"type":"tool_use","id":"toolu_018S9AgPUk5G4hK6NG7KEmLD","name":"Bash","input":{"command":"grep -r \"window\\|crop\\|slice\\|truncate\" /home/ubuntu/ppi_project --include=\"*.py\" ! -path \"*/ppi-env/*\" ! -path \"*/__pycache__/*\" | head -30"},"caller":{"type":"direct"}}],"stop_reason":"tool_use","stop_sequence":null,"stop_details":null,"usage":{"input_tokens":13825,"cache_creation_input_tokens":15130,"cache_read_input_tokens":26543,"output_tokens":217,"server_tool_use":{"web_search_requests":0,"web_fetch_requests":0},"service_tier":"standard","cache_creation":{"ephemeral_1h_input_tokens":0,"ephemeral_5m_input_tokens":15130},"inference_geo":"not_available","iterations":[{"input_tokens":13825,"output_tokens":217,"cache_read_input_tokens":26543,"cache_creation_input_tokens":15130,"cache_creation":{"ephemeral_5m_input_tokens":15130,"ephemeral_1h_input_tokens":0},"type":"message"}],"speed":"standard"},"diagnostics":null},"requestId":"req_011CbcmzSP4j7vMv6uJsT2EM","attributionAgent":"Explore","type":"assistant","uuid":"813a4b0a-8d99-42ec-aca9-c99ee3f8c271","timestamp":"2026-06-01T15:28:24.496Z","userType":"external","entrypoint":"claude-vscode","cwd":"/home/ubuntu/ppi_project","sessionId":"5a87a0cd-6802-4ac5-a708-29b2bf7ea563","version":"2.1.159","gitBranch":"afdb"} +{"parentUuid":"813a4b0a-8d99-42ec-aca9-c99ee3f8c271","isSidechain":true,"promptId":"ae94540a-045a-4847-b16d-c2358b5c9db6","agentId":"a07e7799acd8be342","type":"user","message":{"role":"user","content":[{"tool_use_id":"toolu_018S9AgPUk5G4hK6NG7KEmLD","type":"tool_result","content":"ugrep: warning: !: No such file or directory\nugrep: warning: */ppi-env/*: No such file or directory\nugrep: warning: !: No such file or directory\nugrep: warning: */__pycache__/*: No such file or directory\nugrep: invalid argument -t TYPES, valid arguments are 'actionscript', 'ada', 'adoc', 'asm', 'asp', 'aspx', 'autoconf', 'automake', 'awk', 'Awk', 'basic', 'batch', 'bison', 'c', 'c++', 'clojure', 'cpp', 'csharp', 'css', 'csv', 'dart', 'Dart', 'delphi', 'elisp', 'elixir', 'erlang', 'fortran', 'gif', 'Gif', 'go', 'groovy', 'gsp', 'haskell', 'html', 'jade', 'java', 'jpeg', 'Jpeg', 'js', 'json', 'jsp', 'julia', 'kotlin', 'less', 'lex', 'lisp', 'lua', 'm4', 'make', 'markdown', 'matlab', 'node', 'Node', 'objc', 'objc++', 'ocaml', 'parrot', 'pascal', 'pdf', 'Pdf', 'perl', 'Perl', 'php', 'Php', 'png', 'Png', 'prolog', 'python', 'Python', 'r', 'rpm', 'Rpm', 'rst', 'rtf', 'Rtf', 'ruby', 'Ruby', 'rust', 'scala', 'scheme', 'seed7', 'shell', 'Shell', 'smalltalk', 'sql', 'svg', 'swift', 'tcl', 'tex', 'text', 'tiff', 'Tiff', 'tt', 'typescript', 'verilog', 'vhdl', 'vim', 'xml', 'Xml', 'yacc', 'yaml', 'zig', and 'list' to show a detailed list of file types\nFor more help on options, try `ugrep --help' or `ugrep --help -t'","is_error":false}]},"uuid":"3e283bfe-1203-42a7-a29f-fa92663a16c3","timestamp":"2026-06-01T15:28:25.875Z","sourceToolAssistantUUID":"813a4b0a-8d99-42ec-aca9-c99ee3f8c271","userType":"external","entrypoint":"claude-vscode","cwd":"/home/ubuntu/ppi_project","sessionId":"5a87a0cd-6802-4ac5-a708-29b2bf7ea563","version":"2.1.159","gitBranch":"afdb"} +{"parentUuid":"3e283bfe-1203-42a7-a29f-fa92663a16c3","isSidechain":true,"agentId":"a07e7799acd8be342","message":{"model":"claude-haiku-4-5-20251001","id":"msg_01LaqBrqxQ3ijpcbTgfbHHAo","type":"message","role":"assistant","content":[{"type":"text","text":"Let me fix the grep search:"}],"stop_reason":null,"stop_sequence":null,"stop_details":null,"usage":{"input_tokens":6,"cache_creation_input_tokens":15711,"cache_read_input_tokens":41673,"cache_creation":{"ephemeral_5m_input_tokens":15711,"ephemeral_1h_input_tokens":0},"output_tokens":3,"service_tier":"standard","inference_geo":"not_available"},"diagnostics":null},"requestId":"req_011CbcmzjjLrtLCDQvkt4M1Z","attributionAgent":"Explore","type":"assistant","uuid":"72712236-c463-45ac-912b-a71b4d757d8c","timestamp":"2026-06-01T15:28:27.574Z","userType":"external","entrypoint":"claude-vscode","cwd":"/home/ubuntu/ppi_project","sessionId":"5a87a0cd-6802-4ac5-a708-29b2bf7ea563","version":"2.1.159","gitBranch":"afdb"} +{"parentUuid":"72712236-c463-45ac-912b-a71b4d757d8c","isSidechain":true,"agentId":"a07e7799acd8be342","message":{"model":"claude-haiku-4-5-20251001","id":"msg_01LaqBrqxQ3ijpcbTgfbHHAo","type":"message","role":"assistant","content":[{"type":"tool_use","id":"toolu_017QVMaJzoo8Y7cWeWk2ag5h","name":"Bash","input":{"command":"grep -r \"window\\|crop\\|slice\" /home/ubuntu/ppi_project/*.py 2>/dev/null | head -20"},"caller":{"type":"direct"}}],"stop_reason":null,"stop_sequence":null,"stop_details":null,"usage":{"input_tokens":6,"cache_creation_input_tokens":15711,"cache_read_input_tokens":41673,"cache_creation":{"ephemeral_5m_input_tokens":15711,"ephemeral_1h_input_tokens":0},"output_tokens":3,"service_tier":"standard","inference_geo":"not_available"},"diagnostics":null},"requestId":"req_011CbcmzjjLrtLCDQvkt4M1Z","attributionAgent":"Explore","type":"assistant","uuid":"1faad4a5-e7b2-4cef-bced-cf34530b5eae","timestamp":"2026-06-01T15:28:28.026Z","userType":"external","entrypoint":"claude-vscode","cwd":"/home/ubuntu/ppi_project","sessionId":"5a87a0cd-6802-4ac5-a708-29b2bf7ea563","version":"2.1.159","gitBranch":"afdb"} +{"parentUuid":"1faad4a5-e7b2-4cef-bced-cf34530b5eae","isSidechain":true,"promptId":"ae94540a-045a-4847-b16d-c2358b5c9db6","agentId":"a07e7799acd8be342","type":"user","message":{"role":"user","content":[{"tool_use_id":"toolu_017QVMaJzoo8Y7cWeWk2ag5h","type":"tool_result","content":"/home/ubuntu/ppi_project/evaluate_pdb_benchmark.py: adds_special_tokens: If True, slice off CLS/SEP tokens from masks\n/home/ubuntu/ppi_project/evaluate_pdb_benchmark.py: seq1_mask_sliced = seq1_masks[:, 1:-1]\n/home/ubuntu/ppi_project/evaluate_pdb_benchmark.py: seq2_mask_sliced = seq2_masks[:, 1:-1]\n/home/ubuntu/ppi_project/evaluate_pdb_benchmark.py: seq1_mask_sliced = seq1_masks\n/home/ubuntu/ppi_project/evaluate_pdb_benchmark.py: seq2_mask_sliced = seq2_masks\n/home/ubuntu/ppi_project/evaluate_pdb_benchmark.py: valid_mask = (seq1_mask_sliced.unsqueeze(2) * seq2_mask_sliced.unsqueeze(1)).bool()\n/home/ubuntu/ppi_project/dataset.py: # The target map size matches model output (which slices off CLS/SEP if present)\n/home/ubuntu/ppi_project/eval.py: adds_special_tokens: If True, slice off CLS/SEP tokens from masks. Set False for gLM2.\n/home/ubuntu/ppi_project/eval.py: seq1_mask_sliced = seq1_mask[:, 1:-1]\n/home/ubuntu/ppi_project/eval.py: seq2_mask_sliced = seq2_mask[:, 1:-1]\n/home/ubuntu/ppi_project/eval.py: seq1_mask_sliced = seq1_mask\n/home/ubuntu/ppi_project/eval.py: seq2_mask_sliced = seq2_mask\n/home/ubuntu/ppi_project/eval.py: valid_mask = (seq1_mask_sliced.unsqueeze(2) * seq2_mask_sliced.unsqueeze(1)).bool()\n/home/ubuntu/ppi_project/evaluate_rf2_benchmark.py: adds_special_tokens: If True, slice off CLS/SEP tokens from masks\n/home/ubuntu/ppi_project/evaluate_rf2_benchmark.py: seq1_mask_sliced = seq1_masks[:, 1:-1]\n/home/ubuntu/ppi_project/evaluate_rf2_benchmark.py: seq2_mask_sliced = seq2_masks[:, 1:-1]\n/home/ubuntu/ppi_project/evaluate_rf2_benchmark.py: seq1_mask_sliced = seq1_masks\n/home/ubuntu/ppi_project/evaluate_rf2_benchmark.py: seq2_mask_sliced = seq2_masks\n/home/ubuntu/ppi_project/evaluate_rf2_benchmark.py: valid_mask = (seq1_mask_sliced.unsqueeze(2) * seq2_mask_sliced.unsqueeze(1)).bool()\n/home/ubuntu/ppi_project/eval_checkpoint.py: # Get actual sequence lengths from sequences (contact map is already sliced)","is_error":false}]},"uuid":"563c2008-da2e-4577-9343-8e1c0c23fab7","timestamp":"2026-06-01T15:28:28.059Z","sourceToolAssistantUUID":"1faad4a5-e7b2-4cef-bced-cf34530b5eae","userType":"external","entrypoint":"claude-vscode","cwd":"/home/ubuntu/ppi_project","sessionId":"5a87a0cd-6802-4ac5-a708-29b2bf7ea563","version":"2.1.159","gitBranch":"afdb"} +{"parentUuid":"563c2008-da2e-4577-9343-8e1c0c23fab7","isSidechain":true,"agentId":"a07e7799acd8be342","message":{"model":"claude-haiku-4-5-20251001","id":"msg_01LaqBrqxQ3ijpcbTgfbHHAo","type":"message","role":"assistant","content":[{"type":"tool_use","id":"toolu_01Vhavpyy35RsDJ4sgj9yzqu","name":"Read","input":{"file_path":"/home/ubuntu/ppi_project/config.py"},"caller":{"type":"direct"}}],"stop_reason":"tool_use","stop_sequence":null,"stop_details":null,"usage":{"input_tokens":6,"cache_creation_input_tokens":15711,"cache_read_input_tokens":41673,"output_tokens":140,"server_tool_use":{"web_search_requests":0,"web_fetch_requests":0},"service_tier":"standard","cache_creation":{"ephemeral_1h_input_tokens":0,"ephemeral_5m_input_tokens":15711},"inference_geo":"not_available","iterations":[{"input_tokens":6,"output_tokens":140,"cache_read_input_tokens":41673,"cache_creation_input_tokens":15711,"cache_creation":{"ephemeral_5m_input_tokens":15711,"ephemeral_1h_input_tokens":0},"type":"message"}],"speed":"standard"},"diagnostics":null},"requestId":"req_011CbcmzjjLrtLCDQvkt4M1Z","attributionAgent":"Explore","type":"assistant","uuid":"e35d23e5-ab8a-4282-aa54-caf7decda5ed","timestamp":"2026-06-01T15:28:28.157Z","userType":"external","entrypoint":"claude-vscode","cwd":"/home/ubuntu/ppi_project","sessionId":"5a87a0cd-6802-4ac5-a708-29b2bf7ea563","version":"2.1.159","gitBranch":"afdb"} +{"parentUuid":"e35d23e5-ab8a-4282-aa54-caf7decda5ed","isSidechain":true,"promptId":"ae94540a-045a-4847-b16d-c2358b5c9db6","agentId":"a07e7799acd8be342","type":"user","message":{"role":"user","content":[{"tool_use_id":"toolu_01Vhavpyy35RsDJ4sgj9yzqu","type":"tool_result","content":"1\tfrom dataclasses import dataclass, field\n2\tfrom typing import Optional, Literal\n3\t\n4\t@dataclass\n5\tclass ModelConfig:\n6\t \"\"\"Configuration for the PPI-CLIP model architecture.\"\"\"\n7\t \n8\t plm_model_name: str = field(default=\"facebook/esm2_t6_8M_UR50D\")\n9\t clip_embed_dim: int = field(default=256)\n10\t contact_embed_dim: int = field(default=256)\n11\t max_len: int = field(default=512)\n12\t pooling_method: Literal[\"mean\", \"max\"] = field(default=\"mean\")\n13\t num_frozen_layers: int = field(default=0)\n14\t model_type: Literal[\"clip\", \"late_retrieval\"] = field(default=\"clip\") # \"clip\" for ContrastivePPIModel, \"late_retrieval\" for LateRetrievalPPIModel\n15\t num_latents: int = field(default=16) # Number of latent vectors for LatentPooling in ColBERT model\n16\t gradient_checkpointing: bool = field(default=False) # Enable gradient checkpointing to reduce memory usage\n17\t\n18\t@dataclass\n19\tclass TrainingConfig:\n20\t \"\"\"Configuration for the training process.\"\"\"\n21\t \n22\t # --- Data ---\n23\t dataset_path: str = field()\n24\t train_split: str = field(default=\"train\")\n25\t val_split: str = field(default=\"validation\")\n26\t ddi_dataset: Optional[str] = field(default=None) # Optional DDI dataset to concatenate with train dataset\n27\t afdb_dataset: Optional[str] = field(default=None) # Optional AlphaFold heterodimer dataset to concatenate with train dataset\n28\t neg_dataset: Optional[str] = field(default=None) # Optional negative interaction dataset (no contacts)\n29\t ddi_only: bool = field(default=False) # If True, ignore train split from dataset_path and only use ddi_dataset\n30\t swap_self_negative: bool = field(default=False) # If True, project each pair through both (q,k) and (k,q) to create 2x batch with self-retrieval negatives\n31\t \n32\t # --- Training ---\n33\t batch_size: int = field(default=256) # This is now the LARGE effective batch size\n34\t grad_cache_chunk_size: Optional[int] = field(default=None) # This is the small physical batch size. Set to None to disable grad_cache and use standard loss.\n35\t eval_batch_size: int = field(default=8)\n36\t max_epochs: int = field(default=10)\n37\t learning_rate: float = field(default=1e-4)\n38\t weight_decay: float = field(default=0.01)\n39\t enable_clip: bool = field(default=True)\n40\t enable_contact_prediction: bool = field(default=True)\n41\t clip_loss_weight: float = field(default=1.0)\n42\t contact_loss_weight: float = field(default=1.0)\n43\t contact_random_neg_ratio: float = field(default=1.0) # Multiplier for random negatives\n44\t contact_self_neg_ratio: float = field(default=1.0) # Multiplier for self negatives (A, A)\n45\t contact_hard_neg_ratio: float = field(default=0.0) # Multiplier for hard negatives (most similar by CLIP)\n46\t use_focal_loss: bool = field(default=False) # Use focal loss instead of standard BCE for contact prediction\n47\t warmup_fraction: float = field(default=0.05)\n48\t max_grad_norm: Optional[float] = field(default=1.0)\n49\t seed: int = field(default=42)\n50\t \n51\t # --- Logging & Saving ---\n52\t output_dir: str = field(default=\"ppi_clip_checkpoint\")\n53\t project_name: str = field(default=\"ppi_clip\")\n54\t run_name: Optional[str] = field(default=None)\n55\t use_wandb: bool = field(default=True)\n56\t log_interval: int = field(default=20)\n57\t eval_interval: int = field(default=100)\n58\t save_interval: int = field(default=500)\n59\t max_eval_examples: int = field(default=1000)\n60\t dataloader_num_workers: int = field(default=4)\n61\t compute_contact_metrics: bool = field(default=True) # If True, compute contact prediction metrics during evaluation\n62\t contact_metrics_max_examples: int = field(default=100) # Maximum number of examples to use for contact metrics computation\n63\t retrieval_eval_path: Optional[str] = field(default=None) # Path to retrieval evaluation dataset to local disk\n64\t benchmark_eval_path: Optional[str] = field(default=None) # Path to benchmark evaluation dataset (TSV file or HuggingFace dataset name)\n65\t use_cluster_weighted_sampling: bool = field(default=False) # Use cluster-weighted sampling (sample cluster uniformly, then sample example from cluster)\n66\t ddi_sampling_weight: float = field(default=0.5) # Fraction of steps drawn from the DDI pool\n67\t afdb_sampling_weight: float = field(default=0.0) # Fraction of steps drawn from the AFDB pool; remainder goes to PPI\n68\t filter_homomers: bool = field(default=False) # If True, filter out homomers (uniprot1_id == uniprot2_id OR seq1_cluster == seq2_cluster). DDI datasets are not filtered.\n69\t resume_from_checkpoint: Optional[str] = field(default=None) # Path to accelerate checkpoint to resume from (loads full training state)\n70\t load_from_checkpoint: Optional[str] = field(default=None) # Path to checkpoint to load model weights only (no optimizer/scheduler state)"}]},"uuid":"e7356cc7-27ab-4989-8204-7a46f2425b2f","timestamp":"2026-06-01T15:28:28.566Z","sourceToolAssistantUUID":"e35d23e5-ab8a-4282-aa54-caf7decda5ed","userType":"external","entrypoint":"claude-vscode","cwd":"/home/ubuntu/ppi_project","sessionId":"5a87a0cd-6802-4ac5-a708-29b2bf7ea563","version":"2.1.159","gitBranch":"afdb"} +{"parentUuid":"e7356cc7-27ab-4989-8204-7a46f2425b2f","isSidechain":true,"agentId":"a07e7799acd8be342","message":{"model":"claude-haiku-4-5-20251001","id":"msg_01JQ9fwawbUgyE896UCbikPb","type":"message","role":"assistant","content":[{"type":"text","text":"Now let me check the loss.py file to understand how contact targets are used:"}],"stop_reason":null,"stop_sequence":null,"stop_details":null,"usage":{"input_tokens":6,"cache_creation_input_tokens":2533,"cache_read_input_tokens":57384,"cache_creation":{"ephemeral_5m_input_tokens":2533,"ephemeral_1h_input_tokens":0},"output_tokens":1,"service_tier":"standard","inference_geo":"not_available"},"diagnostics":null},"requestId":"req_011CbcmzwDejAL2Y5hCypdoF","attributionAgent":"Explore","type":"assistant","uuid":"44a6699e-e0bb-40a4-82ae-c42365790c29","timestamp":"2026-06-01T15:28:30.204Z","userType":"external","entrypoint":"claude-vscode","cwd":"/home/ubuntu/ppi_project","sessionId":"5a87a0cd-6802-4ac5-a708-29b2bf7ea563","version":"2.1.159","gitBranch":"afdb"} +{"parentUuid":"44a6699e-e0bb-40a4-82ae-c42365790c29","isSidechain":true,"agentId":"a07e7799acd8be342","message":{"model":"claude-haiku-4-5-20251001","id":"msg_01JQ9fwawbUgyE896UCbikPb","type":"message","role":"assistant","content":[{"type":"tool_use","id":"toolu_01AqQodfaPwNG7kecxNLCtMW","name":"Read","input":{"file_path":"/home/ubuntu/ppi_project/loss.py","limit":100},"caller":{"type":"direct"}}],"stop_reason":null,"stop_sequence":null,"stop_details":null,"usage":{"input_tokens":6,"cache_creation_input_tokens":2533,"cache_read_input_tokens":57384,"cache_creation":{"ephemeral_5m_input_tokens":2533,"ephemeral_1h_input_tokens":0},"output_tokens":1,"service_tier":"standard","inference_geo":"not_available"},"diagnostics":null},"requestId":"req_011CbcmzwDejAL2Y5hCypdoF","attributionAgent":"Explore","type":"assistant","uuid":"1f493893-cca2-496a-aaf2-edab388ea8d5","timestamp":"2026-06-01T15:28:30.684Z","userType":"external","entrypoint":"claude-vscode","cwd":"/home/ubuntu/ppi_project","sessionId":"5a87a0cd-6802-4ac5-a708-29b2bf7ea563","version":"2.1.159","gitBranch":"afdb"} +{"parentUuid":"1f493893-cca2-496a-aaf2-edab388ea8d5","isSidechain":true,"agentId":"a07e7799acd8be342","message":{"model":"claude-haiku-4-5-20251001","id":"msg_01JQ9fwawbUgyE896UCbikPb","type":"message","role":"assistant","content":[{"type":"tool_use","id":"toolu_01DsSqfJt8LaFaSCCKQDH5AG","name":"Bash","input":{"command":"grep -n \"contact_target\\|contact_logits\\|contact_loss\" /home/ubuntu/ppi_project/loss.py | head -30"},"caller":{"type":"direct"}}],"stop_reason":null,"stop_sequence":null,"stop_details":null,"usage":{"input_tokens":6,"cache_creation_input_tokens":2533,"cache_read_input_tokens":57384,"cache_creation":{"ephemeral_5m_input_tokens":2533,"ephemeral_1h_input_tokens":0},"output_tokens":1,"service_tier":"standard","inference_geo":"not_available"},"diagnostics":null},"requestId":"req_011CbcmzwDejAL2Y5hCypdoF","attributionAgent":"Explore","type":"assistant","uuid":"1d037bb5-eea0-45a7-b1b4-0876c1ff6279","timestamp":"2026-06-01T15:28:30.784Z","userType":"external","entrypoint":"claude-vscode","cwd":"/home/ubuntu/ppi_project","sessionId":"5a87a0cd-6802-4ac5-a708-29b2bf7ea563","version":"2.1.159","gitBranch":"afdb"} +{"parentUuid":"1d037bb5-eea0-45a7-b1b4-0876c1ff6279","isSidechain":true,"promptId":"ae94540a-045a-4847-b16d-c2358b5c9db6","agentId":"a07e7799acd8be342","type":"user","message":{"role":"user","content":[{"tool_use_id":"toolu_01DsSqfJt8LaFaSCCKQDH5AG","type":"tool_result","content":"89: def forward(self, contact_logits, contact_target, valid_mask):\n92: contact_logits: (B', L1, L2) contact prediction logits\n93: contact_target: (B, L1, L2) boolean contact targets\n97: return torch.tensor(0.0, device=contact_logits.device, requires_grad=True)\n99: # Expand contact_target if contact_logits has more pairs (negatives)\n100: if contact_logits.shape[0] > contact_target.shape[0]:\n101: num_neg = contact_logits.shape[0] - contact_target.shape[0]\n102: neg_target = torch.zeros(num_neg, *contact_target.shape[1:], \n103: device=contact_target.device, dtype=contact_target.dtype)\n104: contact_target = torch.cat([contact_target, neg_target], dim=0)\n108: loss = focal_loss(contact_logits, contact_target.float(), valid_mask, gamma=2.0, alpha=0.25)\n112: contact_logits[valid_mask], \n113: contact_target.float()[valid_mask],\n509: def __init__(self, model, chunk_size, clip_loss_fn, contact_loss_fn, contact_weight, accelerator, enable_clip=True, clip_weight=1.0):\n513: self.contact_loss_fn = contact_loss_fn\n523: targets = batch['contact_target']\n630: total_contact_loss = 0.0\n664: c_loss = self.contact_loss_fn(c_logits, chunk['targets'], c_valid_mask)\n666: total_contact_loss += c_loss.item()\n673: # Also ensure contact_logits are attached if they exist but c_loss ended up being 0/detached\n691: return clip_loss_val, total_contact_loss, logit_scale_val, top1_acc","is_error":false}]},"uuid":"17023ffe-e425-4c51-aae8-a6874f72a05c","timestamp":"2026-06-01T15:28:30.809Z","sourceToolAssistantUUID":"1d037bb5-eea0-45a7-b1b4-0876c1ff6279","userType":"external","entrypoint":"claude-vscode","cwd":"/home/ubuntu/ppi_project","sessionId":"5a87a0cd-6802-4ac5-a708-29b2bf7ea563","version":"2.1.159","gitBranch":"afdb"} +{"parentUuid":"1f493893-cca2-496a-aaf2-edab388ea8d5","isSidechain":true,"promptId":"ae94540a-045a-4847-b16d-c2358b5c9db6","agentId":"a07e7799acd8be342","type":"user","message":{"role":"user","content":[{"tool_use_id":"toolu_01AqQodfaPwNG7kecxNLCtMW","type":"tool_result","content":"1\t# loss.py \n2\timport torch\n3\timport torch.nn as nn\n4\timport torch.nn.functional as F\n5\tfrom contextlib import nullcontext # Import needed\n6\t\n7\ttry:\n8\t import torch.distributed.nn\n9\t from torch import distributed as dist\n10\t has_distributed = True\n11\texcept ImportError:\n12\t has_distributed = False\n13\t\n14\t# --- Helper Functions ---\n15\t\n16\tdef maxsim(q_reps, d_reps):\n17\t \"\"\"\n18\t Memory-efficient MaxSim for ColBERT-style embeddings.\n19\t \n20\t Args:\n21\t q_reps: (B_q, K_q, D) - Query latent representations\n22\t d_reps: (B_d, K_d, D) - Document latent representations\n23\t \n24\t Returns:\n25\t scores: (B_q, B_d) - MaxSim similarity scores\n26\t \n27\t Note: No masking needed since LatentPooling outputs fixed K latents (no padding).\n28\t \"\"\"\n29\t B_q, K_q, D = q_reps.shape\n30\t B_d, K_d, _ = d_reps.shape\n31\t \n32\t scores = torch.zeros(B_q, B_d, device=q_reps.device, dtype=q_reps.dtype)\n33\t \n34\t # Flatten all documents: (B_d, K_d, D) -> (D, B_d * K_d)\n35\t d_flat = d_reps.permute(2, 0, 1).reshape(D, -1)\n36\t \n37\t # Chunk over queries only\n38\t chunk_size = 128\n39\t \n40\t for i in range(0, B_q, chunk_size):\n41\t q_chunk = q_reps[i:i+chunk_size] # (chunk, K_q, D)\n42\t \n43\t # Interaction: (chunk, K_q, D) @ (D, B_d * K_d) -> (chunk, K_q, B_d * K_d)\n44\t sim_chunk = q_chunk @ d_flat\n45\t \n46\t # Reshape: (chunk, K_q, B_d, K_d)\n47\t sim_chunk = sim_chunk.view(q_chunk.shape[0], K_q, B_d, K_d)\n48\t \n49\t # Max over document latents: (chunk, K_q, B_d)\n50\t max_sim, _ = sim_chunk.max(dim=-1)\n51\t \n52\t # Sum over query latents: (chunk, B_d)\n53\t chunk_scores = max_sim.sum(dim=1)\n54\t \n55\t scores[i:i+chunk_size] = chunk_scores\n56\t \n57\t return scores\n58\t\n59\t# --- 1. Standard Losses ---\n60\t\n61\tdef focal_loss(logits, targets, valid_mask, gamma=2.0, alpha=0.25):\n62\t \"\"\"\n63\t Binary focal loss.\n64\t Normalized by the number of positive examples to prevent gradient vanishing\n65\t in sparse contact maps.\n66\t \"\"\"\n67\t if valid_mask.sum() == 0:\n68\t return torch.tensor(0.0, device=logits.device, requires_grad=True)\n69\t \n70\t logits_valid = logits[valid_mask]\n71\t targets_valid = targets[valid_mask]\n72\t \n73\t bce_loss = F.binary_cross_entropy_with_logits(logits_valid, targets_valid, reduction='none')\n74\t pt = torch.exp(-bce_loss)\n75\t alpha_tensor = (1 - alpha) + targets_valid * (2 * alpha - 1)\n76\t f_loss = alpha_tensor * (1 - pt) ** gamma * bce_loss\n77\t \n78\t num_positives = targets_valid.sum()\n79\t if num_positives == 0:\n80\t return f_loss.sum()\n81\t else:\n82\t return f_loss.sum() / num_positives\n83\t\n84\tclass ContactLoss(nn.Module):\n85\t def __init__(self, use_focal_loss=False):\n86\t super().__init__()\n87\t self.use_focal_loss = use_focal_loss\n88\t\n89\t def forward(self, contact_logits, contact_target, valid_mask):\n90\t \"\"\"\n91\t Args:\n92\t contact_logits: (B', L1, L2) contact prediction logits\n93\t contact_target: (B, L1, L2) boolean contact targets\n94\t valid_mask: (B', L1, L2) boolean mask, True for valid positions, False for padding\n95\t \"\"\"\n96\t if valid_mask.sum() == 0:\n97\t return torch.tensor(0.0, device=contact_logits.device, requires_grad=True)\n98\t \n99\t # Expand contact_target if contact_logits has more pairs (negatives)\n100\t if contact_logits.shape[0] > contact_target.shape[0]:"}]},"uuid":"f7e90e7f-c6f1-42e4-82df-fadfd9035642","timestamp":"2026-06-01T15:28:31.198Z","sourceToolAssistantUUID":"1f493893-cca2-496a-aaf2-edab388ea8d5","userType":"external","entrypoint":"claude-vscode","cwd":"/home/ubuntu/ppi_project","sessionId":"5a87a0cd-6802-4ac5-a708-29b2bf7ea563","version":"2.1.159","gitBranch":"afdb"} +{"parentUuid":"f7e90e7f-c6f1-42e4-82df-fadfd9035642","isSidechain":true,"agentId":"a07e7799acd8be342","message":{"model":"claude-haiku-4-5-20251001","id":"msg_01LauGgjZjTWkYdEsWSUmsmV","type":"message","role":"assistant","content":[{"type":"text","text":"Now let me read more of the loss file to understand the contact loss better:"}],"stop_reason":null,"stop_sequence":null,"stop_details":null,"usage":{"input_tokens":527,"cache_creation_input_tokens":1557,"cache_read_input_tokens":59917,"cache_creation":{"ephemeral_5m_input_tokens":1557,"ephemeral_1h_input_tokens":0},"output_tokens":2,"service_tier":"standard","inference_geo":"not_available"},"diagnostics":null},"requestId":"req_011Cbcn18WZBXZBLz2rjgWKc","attributionAgent":"Explore","type":"assistant","uuid":"28bdcd82-0a2b-4d39-a6d1-3e3a40886468","timestamp":"2026-06-01T15:28:32.639Z","userType":"external","entrypoint":"claude-vscode","cwd":"/home/ubuntu/ppi_project","sessionId":"5a87a0cd-6802-4ac5-a708-29b2bf7ea563","version":"2.1.159","gitBranch":"afdb"} +{"parentUuid":"28bdcd82-0a2b-4d39-a6d1-3e3a40886468","isSidechain":true,"agentId":"a07e7799acd8be342","message":{"model":"claude-haiku-4-5-20251001","id":"msg_01LauGgjZjTWkYdEsWSUmsmV","type":"message","role":"assistant","content":[{"type":"tool_use","id":"toolu_01NzYbqoASC75T9PVyEHzSt3","name":"Read","input":{"file_path":"/home/ubuntu/ppi_project/loss.py","offset":84,"limit":50},"caller":{"type":"direct"}}],"stop_reason":"tool_use","stop_sequence":null,"stop_details":null,"usage":{"input_tokens":527,"cache_creation_input_tokens":1557,"cache_read_input_tokens":59917,"output_tokens":116,"server_tool_use":{"web_search_requests":0,"web_fetch_requests":0},"service_tier":"standard","cache_creation":{"ephemeral_1h_input_tokens":0,"ephemeral_5m_input_tokens":1557},"inference_geo":"not_available","iterations":[{"input_tokens":527,"output_tokens":116,"cache_read_input_tokens":59917,"cache_creation_input_tokens":1557,"cache_creation":{"ephemeral_5m_input_tokens":1557,"ephemeral_1h_input_tokens":0},"type":"message"}],"speed":"standard"},"diagnostics":null},"requestId":"req_011Cbcn18WZBXZBLz2rjgWKc","attributionAgent":"Explore","type":"assistant","uuid":"a91ef84e-0761-4208-b1cc-9b6c4c162e9f","timestamp":"2026-06-01T15:28:33.043Z","userType":"external","entrypoint":"claude-vscode","cwd":"/home/ubuntu/ppi_project","sessionId":"5a87a0cd-6802-4ac5-a708-29b2bf7ea563","version":"2.1.159","gitBranch":"afdb"} +{"parentUuid":"a91ef84e-0761-4208-b1cc-9b6c4c162e9f","isSidechain":true,"promptId":"ae94540a-045a-4847-b16d-c2358b5c9db6","agentId":"a07e7799acd8be342","type":"user","message":{"role":"user","content":[{"tool_use_id":"toolu_01NzYbqoASC75T9PVyEHzSt3","type":"tool_result","content":"84\tclass ContactLoss(nn.Module):\n85\t def __init__(self, use_focal_loss=False):\n86\t super().__init__()\n87\t self.use_focal_loss = use_focal_loss\n88\t\n89\t def forward(self, contact_logits, contact_target, valid_mask):\n90\t \"\"\"\n91\t Args:\n92\t contact_logits: (B', L1, L2) contact prediction logits\n93\t contact_target: (B, L1, L2) boolean contact targets\n94\t valid_mask: (B', L1, L2) boolean mask, True for valid positions, False for padding\n95\t \"\"\"\n96\t if valid_mask.sum() == 0:\n97\t return torch.tensor(0.0, device=contact_logits.device, requires_grad=True)\n98\t \n99\t # Expand contact_target if contact_logits has more pairs (negatives)\n100\t if contact_logits.shape[0] > contact_target.shape[0]:\n101\t num_neg = contact_logits.shape[0] - contact_target.shape[0]\n102\t neg_target = torch.zeros(num_neg, *contact_target.shape[1:], \n103\t device=contact_target.device, dtype=contact_target.dtype)\n104\t contact_target = torch.cat([contact_target, neg_target], dim=0)\n105\t\n106\t # Cast boolean targets to float for loss computation\n107\t if self.use_focal_loss:\n108\t loss = focal_loss(contact_logits, contact_target.float(), valid_mask, gamma=2.0, alpha=0.25)\n109\t else:\n110\t # Only compute loss on valid positions\n111\t loss = F.binary_cross_entropy_with_logits(\n112\t contact_logits[valid_mask], \n113\t contact_target.float()[valid_mask],\n114\t reduction='mean'\n115\t )\n116\t \n117\t return loss\n118\t\n119\tclass ClipLoss(nn.Module):\n120\t \"\"\"Standard InfoNCE for Evaluation (No GradCache logic).\"\"\"\n121\t def __init__(self, local_loss=False, gather_with_grad=False, rank=0, world_size=1, valid_keys_tensor=None, multiplier=1_000_000):\n122\t super().__init__()\n123\t self.local_loss = local_loss\n124\t self.gather_with_grad = gather_with_grad\n125\t self.rank = rank\n126\t self.world_size = world_size\n127\t self.multiplier = multiplier\n128\t # Register keys as buffer so they automatically move to GPU with .to(device)\n129\t if valid_keys_tensor is not None:\n130\t self.register_buffer('valid_keys', valid_keys_tensor)\n131\t else:\n132\t self.register_buffer('valid_keys', None)\n133\t"}]},"uuid":"eb6ef0d3-f6d5-47d5-acc0-fab227b67900","timestamp":"2026-06-01T15:28:33.053Z","sourceToolAssistantUUID":"a91ef84e-0761-4208-b1cc-9b6c4c162e9f","userType":"external","entrypoint":"claude-vscode","cwd":"/home/ubuntu/ppi_project","sessionId":"5a87a0cd-6802-4ac5-a708-29b2bf7ea563","version":"2.1.159","gitBranch":"afdb"} +{"parentUuid":"eb6ef0d3-f6d5-47d5-acc0-fab227b67900","isSidechain":true,"agentId":"a07e7799acd8be342","message":{"model":"claude-haiku-4-5-20251001","id":"msg_01X8PsduCwn6pK3eHkK9Q4wn","type":"message","role":"assistant","content":[{"type":"text","text":"Now let me check for any eval.py to understand contact evaluation:"}],"stop_reason":null,"stop_sequence":null,"stop_details":null,"usage":{"input_tokens":5,"cache_creation_input_tokens":1401,"cache_read_input_tokens":61474,"cache_creation":{"ephemeral_5m_input_tokens":1401,"ephemeral_1h_input_tokens":0},"output_tokens":1,"service_tier":"standard","inference_geo":"not_available"},"diagnostics":null},"requestId":"req_011Cbcn1Gggs2JxhFJ34wGTP","attributionAgent":"Explore","type":"assistant","uuid":"b2d70faf-85bd-4d5b-ae27-a5d65b3e23fa","timestamp":"2026-06-01T15:28:34.440Z","userType":"external","entrypoint":"claude-vscode","cwd":"/home/ubuntu/ppi_project","sessionId":"5a87a0cd-6802-4ac5-a708-29b2bf7ea563","version":"2.1.159","gitBranch":"afdb"} +{"parentUuid":"b2d70faf-85bd-4d5b-ae27-a5d65b3e23fa","isSidechain":true,"agentId":"a07e7799acd8be342","message":{"model":"claude-haiku-4-5-20251001","id":"msg_01X8PsduCwn6pK3eHkK9Q4wn","type":"message","role":"assistant","content":[{"type":"tool_use","id":"toolu_01De8gsyysLKiBQjPN6mQ93F","name":"Read","input":{"file_path":"/home/ubuntu/ppi_project/eval.py","limit":150},"caller":{"type":"direct"}}],"stop_reason":"tool_use","stop_sequence":null,"stop_details":null,"usage":{"input_tokens":5,"cache_creation_input_tokens":1401,"cache_read_input_tokens":61474,"output_tokens":97,"server_tool_use":{"web_search_requests":0,"web_fetch_requests":0},"service_tier":"standard","cache_creation":{"ephemeral_1h_input_tokens":0,"ephemeral_5m_input_tokens":1401},"inference_geo":"not_available","iterations":[{"input_tokens":5,"output_tokens":97,"cache_read_input_tokens":61474,"cache_creation_input_tokens":1401,"cache_creation":{"ephemeral_5m_input_tokens":1401,"ephemeral_1h_input_tokens":0},"type":"message"}],"speed":"standard"},"diagnostics":null},"requestId":"req_011Cbcn1Gggs2JxhFJ34wGTP","attributionAgent":"Explore","type":"assistant","uuid":"1d163873-54c4-4820-b612-0f66444cd94d","timestamp":"2026-06-01T15:28:34.545Z","userType":"external","entrypoint":"claude-vscode","cwd":"/home/ubuntu/ppi_project","sessionId":"5a87a0cd-6802-4ac5-a708-29b2bf7ea563","version":"2.1.159","gitBranch":"afdb"} +{"parentUuid":"1d163873-54c4-4820-b612-0f66444cd94d","isSidechain":true,"promptId":"ae94540a-045a-4847-b16d-c2358b5c9db6","agentId":"a07e7799acd8be342","type":"user","message":{"role":"user","content":[{"tool_use_id":"toolu_01De8gsyysLKiBQjPN6mQ93F","type":"tool_result","content":"1\t\"\"\"\n2\tEvaluation functions for CLIP and contact prediction metrics.\n3\t\"\"\"\n4\timport csv\n5\timport numpy as np\n6\timport torch\n7\timport torch.nn.functional as F\n8\tfrom collections import defaultdict\n9\tfrom pathlib import Path\n10\tfrom typing import Dict, List, Optional, Tuple\n11\tfrom tqdm.auto import tqdm\n12\tfrom accelerate import Accelerator\n13\tfrom accelerate.logging import get_logger\n14\tfrom sklearn.metrics import average_precision_score, roc_auc_score, precision_recall_curve, auc\n15\timport pandas as pd\n16\tfrom transformers import AutoTokenizer\n17\t\n18\tfrom dataset import create_contact_targets\n19\tfrom loss import maxsim\n20\t\n21\t\n22\tlogger = get_logger(__name__)\n23\t\n24\t\n25\tdef add_prefix_to_dict(d: dict, prefix: str) -> dict:\n26\t \"\"\"Add a prefix to all keys in a dictionary.\"\"\"\n27\t return {f\"{prefix}/{k}\": v for k, v in d.items()}\n28\t\n29\tdef get_clip_metrics(seq1_features, seq2_features, logit_scale, seq1_clusters=None, seq2_clusters=None, valid_keys_tensor=None, multiplier=1_000_000):\n30\t \"\"\"\n31\t Compute CLIP retrieval metrics.\n32\t \n33\t Args:\n34\t seq1_features: (N, D) or (N, L, D) - Normalized CLIP embeddings for sequence 1\n35\t seq2_features: (N, D) or (N, L, D) - Normalized CLIP embeddings for sequence 2\n36\t logit_scale: () - Logit scale parameter\n37\t seq1_clusters: Optional (N,) tensor of cluster IDs for sequence 1\n38\t seq2_clusters: Optional (N,) tensor of cluster IDs for sequence 2\n39\t valid_keys_tensor: Optional tensor of valid interaction keys (from build_interaction_keys)\n40\t multiplier: Multiplier for hash key generation (default 1M)\n41\t \n42\t Returns:\n43\t dict: Dictionary of retrieval metrics\n44\t \"\"\"\n45\t metrics = {}\n46\t N = len(seq1_features)\n47\t \n48\t # Create valid match mask using cluster IDs\n49\t if seq1_clusters is not None and seq2_clusters is not None and valid_keys_tensor is not None and len(valid_keys_tensor) > 0:\n50\t # Generate keys for the batch: (N, 1) * M + (1, N) -> (N, N) keys\n51\t batch_keys = seq1_clusters.unsqueeze(1) * multiplier + seq2_clusters.unsqueeze(0)\n52\t \n53\t # Ensure valid_keys is on the same device as batch_keys\n54\t valid_keys = valid_keys_tensor.to(batch_keys.device)\n55\t \n56\t # Check against the Global Truth Map using torch.isin\n57\t # Note: valid_keys already contains both forward and reverse keys (from build_interaction_keys),\n58\t # so we only need to check one direction, matching the loss function behavior\n59\t valid_match = torch.isin(batch_keys, valid_keys)\n60\t \n61\t # The diagonal is always the positive pair in the batch\n62\t valid_match = valid_match | torch.eye(N, dtype=torch.bool, device=valid_match.device)\n63\t else:\n64\t # Fallback: only diagonal is valid (same-index pairs)\n65\t valid_match = torch.eye(N, dtype=torch.bool)\n66\t \n67\t # Compute similarity matrices\n68\t # Using float32 for matrix multiplication precision\n69\t seq1_features = seq1_features.float()\n70\t seq2_features = seq2_features.float()\n71\t logit_scale = logit_scale.float()\n72\t \n73\t # Check if we have ColBERT-style per-residue embeddings (3D) or CLIP-style pooled (2D)\n74\t is_colbert = len(seq1_features.shape) == 3\n75\t \n76\t if is_colbert:\n77\t # ColBERT MaxSim similarity (runs on GPU for speed)\n78\t logits_seq1_to_seq2 = maxsim(seq1_features, seq2_features) * logit_scale\n79\t logits_seq2_to_seq1 = maxsim(seq2_features, seq1_features) * logit_scale\n80\t # Move results back to CPU for ranking/metrics computation\n81\t logits_seq1_to_seq2 = logits_seq1_to_seq2.cpu()\n82\t logits_seq2_to_seq1 = logits_seq2_to_seq1.cpu()\n83\t # Ensure valid_match is also on CPU\n84\t valid_match = valid_match.cpu()\n85\t else:\n86\t # Standard CLIP dot product similarity (On CPU to avoid OOM for large batches)\n87\t logits_seq1_to_seq2 = (logit_scale * seq1_features @ seq2_features.t())\n88\t logits_seq2_to_seq1 = (logit_scale * seq2_features @ seq1_features.t())\n89\t \n90\t logits = {\n91\t \"seq1_to_seq2\": (logits_seq1_to_seq2, valid_match),\n92\t \"seq2_to_seq1\": (logits_seq2_to_seq1, valid_match.t())\n93\t }\n94\t \n95\t for name, (logit, valid_mask) in logits.items():\n96\t # Get rankings (descending order)\n97\t ranking = torch.argsort(logit, descending=True, dim=1) # (N, N)\n98\t \n99\t # Vectorized: find first valid match for each query\n100\t # valid_mask[i, ranking[i, j]] tells us if ranking position j is valid for query i\n101\t valid_in_ranking = torch.gather(valid_mask, 1, ranking) # (N, N)\n102\t # Find first True in each row (argmax returns 0 if all False, so we need to check)\n103\t first_valid_idx = valid_in_ranking.int().argmax(dim=1) # (N,)\n104\t # Check if the found position is actually valid\n105\t is_valid = torch.gather(valid_in_ranking, 1, first_valid_idx.unsqueeze(1)).squeeze(1) # (N,)\n106\t # Rank is 1-indexed, use N+1 if no valid match found\n107\t ranks = torch.where(is_valid, first_valid_idx + 1, torch.tensor(N + 1)).numpy()\n108\t \n109\t metrics[f\"{name}_mean_rank\"] = float(ranks.mean())\n110\t metrics[f\"{name}_median_rank\"] = float(np.floor(np.median(ranks)))\n111\t \n112\t # Recall@K metrics: check if any of top-K are valid matches (vectorized)\n113\t for k in [1, 5, 10]:\n114\t top_k_valid = valid_in_ranking[:, :k].any(dim=1)\n115\t metrics[f\"{name}_R@{k}\"] = float(top_k_valid.float().mean().item())\n116\t \n117\t return metrics\n118\t\n119\t\n120\tdef get_retrieval_metrics_with_equivalents(\n121\t query_embeddings, \n122\t corpus_embeddings, \n123\t queries, \n124\t corpus, \n125\t logit_scale,\n126\t accept_all_binders=False,\n127\t get_per_query_results=False,\n128\t mask_self_hits=True,\n129\t):\n130\t \"\"\"\n131\t Evaluate retrieval with sequence-level matching.\n132\t \n133\t Accepts matches to:\n134\t 1. Exact binder_uniprot_id match (always valid)\n135\t 2. Any binder in all_binder_uniprot_ids (if accept_all_binders=True): accepts any protein the query actually binds to\n136\t \n137\t Args:\n138\t query_embeddings: (N_q, D) or (N_q, L, D) - Query embeddings\n139\t corpus_embeddings: (N_c, D) or (N_c, L, D) - Corpus embeddings\n140\t queries: List of query dicts with 'binder_uniprot_id' and optionally 'all_binder_uniprot_ids'\n141\t corpus: List of corpus dicts with 'uniprot_id'\n142\t logit_scale: Logit scale parameter\n143\t accept_all_binders: If True, accept matches to any protein in query's all_binder_uniprot_ids list\n144\t \n145\t Returns:\n146\t dict: Dictionary of retrieval metrics\n147\t \"\"\" \n148\t N_q = len(query_embeddings)\n149\t N_c = len(corpus_embeddings)\n150\t "}]},"uuid":"5194c01a-deb9-4f86-bbe4-1cdec7f23ec6","timestamp":"2026-06-01T15:28:34.975Z","sourceToolAssistantUUID":"1d163873-54c4-4820-b612-0f66444cd94d","userType":"external","entrypoint":"claude-vscode","cwd":"/home/ubuntu/ppi_project","sessionId":"5a87a0cd-6802-4ac5-a708-29b2bf7ea563","version":"2.1.159","gitBranch":"afdb"} +{"parentUuid":"5194c01a-deb9-4f86-bbe4-1cdec7f23ec6","isSidechain":true,"agentId":"a07e7799acd8be342","message":{"model":"claude-haiku-4-5-20251001","id":"msg_01JCwuPCqTfHUG6HcxpRT3mw","type":"message","role":"assistant","content":[{"type":"text","text":"Let me also check layers.py for understanding the contact head and architecture:"}],"stop_reason":null,"stop_sequence":null,"stop_details":null,"usage":{"input_tokens":5,"cache_creation_input_tokens":2465,"cache_read_input_tokens":62875,"cache_creation":{"ephemeral_5m_input_tokens":2465,"ephemeral_1h_input_tokens":0},"output_tokens":3,"service_tier":"standard","inference_geo":"not_available"},"diagnostics":null},"requestId":"req_011Cbcn1QgPu2F7NPTCZCq6G","attributionAgent":"Explore","type":"assistant","uuid":"07fa8a71-2ca1-429f-8b6e-ac6e0ec04e04","timestamp":"2026-06-01T15:28:36.684Z","userType":"external","entrypoint":"claude-vscode","cwd":"/home/ubuntu/ppi_project","sessionId":"5a87a0cd-6802-4ac5-a708-29b2bf7ea563","version":"2.1.159","gitBranch":"afdb"} +{"parentUuid":"07fa8a71-2ca1-429f-8b6e-ac6e0ec04e04","isSidechain":true,"agentId":"a07e7799acd8be342","message":{"model":"claude-haiku-4-5-20251001","id":"msg_01JCwuPCqTfHUG6HcxpRT3mw","type":"message","role":"assistant","content":[{"type":"tool_use","id":"toolu_018c9JwG6Q7ywm4w3QCwPcsZ","name":"Read","input":{"file_path":"/home/ubuntu/ppi_project/layers.py","limit":100},"caller":{"type":"direct"}}],"stop_reason":"tool_use","stop_sequence":null,"stop_details":null,"usage":{"input_tokens":5,"cache_creation_input_tokens":2465,"cache_read_input_tokens":62875,"output_tokens":98,"server_tool_use":{"web_search_requests":0,"web_fetch_requests":0},"service_tier":"standard","cache_creation":{"ephemeral_1h_input_tokens":0,"ephemeral_5m_input_tokens":2465},"inference_geo":"not_available","iterations":[{"input_tokens":5,"output_tokens":98,"cache_read_input_tokens":62875,"cache_creation_input_tokens":2465,"cache_creation":{"ephemeral_5m_input_tokens":2465,"ephemeral_1h_input_tokens":0},"type":"message"}],"speed":"standard"},"diagnostics":null},"requestId":"req_011Cbcn1QgPu2F7NPTCZCq6G","attributionAgent":"Explore","type":"assistant","uuid":"18da7f3a-65fe-47ab-be51-d948c47f25ab","timestamp":"2026-06-01T15:28:37.108Z","userType":"external","entrypoint":"claude-vscode","cwd":"/home/ubuntu/ppi_project","sessionId":"5a87a0cd-6802-4ac5-a708-29b2bf7ea563","version":"2.1.159","gitBranch":"afdb"} +{"parentUuid":"18da7f3a-65fe-47ab-be51-d948c47f25ab","isSidechain":true,"promptId":"ae94540a-045a-4847-b16d-c2358b5c9db6","agentId":"a07e7799acd8be342","type":"user","message":{"role":"user","content":[{"tool_use_id":"toolu_018c9JwG6Q7ywm4w3QCwPcsZ","type":"tool_result","content":"1\timport torch\n2\timport torch.nn as nn\n3\timport torch.nn.functional as F\n4\tfrom einops import rearrange, repeat\n5\tfrom typing import Optional, Tuple, Union\n6\tfrom torch.utils.checkpoint import checkpoint\n7\t\n8\ttry:\n9\t from flash_attn.ops.activations import swiglu\n10\t from flash_attn.layers.rotary import apply_rotary_emb_func\n11\t from flash_attn import flash_attn_varlen_kvpacked_func\n12\t from flash_attn.bert_padding import pad_input, unpad_input\n13\t from flash_attn.ops.triton.layer_norm import RMSNorm\n14\t\n15\t FLASH_ATTN_AVAILABLE = True\n16\texcept ImportError:\n17\t print(\"Flash Attention not available, using PyTorch SDPA fallback\")\n18\t FLASH_ATTN_AVAILABLE = False\n19\t unpad_input = pad_input = apply_rotary_emb_func = None\n20\t flash_attn_varlen_kvpacked_func = None\n21\t\n22\t def swiglu(x, y):\n23\t return F.silu(x) * y\n24\t\n25\t def rmsnorm_func(hidden_states, weight, variance_epsilon):\n26\t \"\"\"Apply the root mean square normalization.\"\"\"\n27\t input_dtype = hidden_states.dtype\n28\t hidden_states = hidden_states.to(torch.float32)\n29\t variance = hidden_states.pow(2).mean(-1, keepdim=True)\n30\t hidden_states = hidden_states * torch.rsqrt(variance + variance_epsilon)\n31\t return (weight * hidden_states).to(input_dtype)\n32\t\n33\t class RMSNorm(nn.Module):\n34\t \"\"\"Root mean square normalization.\"\"\"\n35\t\n36\t def __init__(self, dim, eps=1e-6):\n37\t super().__init__()\n38\t self.weight = nn.Parameter(torch.ones(dim))\n39\t self.register_buffer(\n40\t \"variance_epsilon\",\n41\t torch.tensor(eps),\n42\t persistent=False,\n43\t )\n44\t\n45\t def forward(self, hidden_states):\n46\t return rmsnorm_func(hidden_states, self.weight, self.variance_epsilon)\n47\t\n48\t\n49\tclass MLPHead(nn.Module):\n50\t \"\"\"SwiGLU MLP head.\"\"\"\n51\t\n52\t def __init__(self, in_dim: int, out_dim: int, hidden_mult: float = 2.0):\n53\t super().__init__()\n54\t hidden_dim = int(in_dim * hidden_mult)\n55\t self.w1 = nn.Linear(in_dim, hidden_dim, bias=False)\n56\t self.w2 = nn.Linear(hidden_dim, out_dim, bias=False)\n57\t self.w3 = nn.Linear(in_dim, hidden_dim, bias=False)\n58\t\n59\t def forward(self, x):\n60\t return self.w2(F.silu(self.w1(x)) * self.w3(x))\n61\t\n62\t\n63\tclass PairformerConfig:\n64\t def __init__(self, dim, heads, depth, out_dim, ffn_mult=None, use_rope=True):\n65\t self.dim = dim\n66\t self.heads = heads\n67\t self.depth = depth\n68\t self.swiglu_multiple_of = 256\n69\t self.ffn_dim_multiplier = ffn_mult\n70\t self.out_dim = out_dim\n71\t self.norm_eps = 1e-6\n72\t self.use_rope = use_rope\n73\t\n74\t\n75\tdef rotate_half(x, interleaved=False):\n76\t if not interleaved:\n77\t x1, x2 = x.chunk(2, dim=-1)\n78\t return torch.cat((-x2, x1), dim=-1)\n79\t else:\n80\t x1, x2 = x[..., ::2], x[..., 1::2]\n81\t return rearrange(\n82\t torch.stack((-x2, x1), dim=-1), \"... d two -> ... (d two)\", two=2\n83\t )\n84\t\n85\t\n86\tdef apply_rotary_emb_torch(x, cos, sin, interleaved=False, position_ids=None):\n87\t \"\"\"\n88\t x: (batch_size, seqlen, nheads, headdim)\n89\t cos, sin: (seqlen, rotary_dim / 2)\n90\t position_ids: (batch_size, seqlen) - OPTIONAL, used to gather indices\n91\t \"\"\"\n92\t if position_ids is not None:\n93\t cos = cos[position_ids]\n94\t sin = sin[position_ids]\n95\t else:\n96\t cos = cos[: x.shape[1]]\n97\t sin = sin[: x.shape[1]]\n98\t\n99\t if not interleaved:\n100\t cos = repeat(cos, \"... d -> ... 1 (2 d)\")"}]},"uuid":"ae3a486e-fda9-4ce8-a1fe-a8174eeb60b9","timestamp":"2026-06-01T15:28:37.534Z","sourceToolAssistantUUID":"18da7f3a-65fe-47ab-be51-d948c47f25ab","userType":"external","entrypoint":"claude-vscode","cwd":"/home/ubuntu/ppi_project","sessionId":"5a87a0cd-6802-4ac5-a708-29b2bf7ea563","version":"2.1.159","gitBranch":"afdb"} +{"parentUuid":"ae3a486e-fda9-4ce8-a1fe-a8174eeb60b9","isSidechain":true,"agentId":"a07e7799acd8be342","message":{"model":"claude-haiku-4-5-20251001","id":"msg_012TerebKZjL7rPtNTr48sYF","type":"message","role":"assistant","content":[{"type":"text","text":"Now let me search for any dataset loading related files and check what fields samples have:"}],"stop_reason":null,"stop_sequence":null,"stop_details":null,"usage":{"input_tokens":5,"cache_creation_input_tokens":1466,"cache_read_input_tokens":65340,"cache_creation":{"ephemeral_5m_input_tokens":1466,"ephemeral_1h_input_tokens":0},"output_tokens":1,"service_tier":"standard","inference_geo":"not_available"},"diagnostics":null},"requestId":"req_011Cbcn1bbUYcACR8oN2Co2F","attributionAgent":"Explore","type":"assistant","uuid":"5ef8f230-9c8e-48d4-98a5-bd490378bcc5","timestamp":"2026-06-01T15:28:38.911Z","userType":"external","entrypoint":"claude-vscode","cwd":"/home/ubuntu/ppi_project","sessionId":"5a87a0cd-6802-4ac5-a708-29b2bf7ea563","version":"2.1.159","gitBranch":"afdb"} +{"parentUuid":"5ef8f230-9c8e-48d4-98a5-bd490378bcc5","isSidechain":true,"agentId":"a07e7799acd8be342","message":{"model":"claude-haiku-4-5-20251001","id":"msg_012TerebKZjL7rPtNTr48sYF","type":"message","role":"assistant","content":[{"type":"tool_use","id":"toolu_01R7QDVTJSeWgcktJJ8axMcH","name":"Bash","input":{"command":"grep -n \"seq1\\|seq2\\|contact\\|cluster\" /home/ubuntu/ppi_project/dataset.py | head -50"},"caller":{"type":"direct"}}],"stop_reason":"tool_use","stop_sequence":null,"stop_details":null,"usage":{"input_tokens":5,"cache_creation_input_tokens":1466,"cache_read_input_tokens":65340,"output_tokens":103,"server_tool_use":{"web_search_requests":0,"web_fetch_requests":0},"service_tier":"standard","cache_creation":{"ephemeral_1h_input_tokens":0,"ephemeral_5m_input_tokens":1466},"inference_geo":"not_available","iterations":[{"input_tokens":5,"output_tokens":103,"cache_read_input_tokens":65340,"cache_creation_input_tokens":1466,"cache_creation":{"ephemeral_5m_input_tokens":1466,"ephemeral_1h_input_tokens":0},"type":"message"}],"speed":"standard"},"diagnostics":null},"requestId":"req_011Cbcn1bbUYcACR8oN2Co2F","attributionAgent":"Explore","type":"assistant","uuid":"e3920087-d844-46fa-b9cc-916f22123c21","timestamp":"2026-06-01T15:28:39.401Z","userType":"external","entrypoint":"claude-vscode","cwd":"/home/ubuntu/ppi_project","sessionId":"5a87a0cd-6802-4ac5-a708-29b2bf7ea563","version":"2.1.159","gitBranch":"afdb"} +{"parentUuid":"e3920087-d844-46fa-b9cc-916f22123c21","isSidechain":true,"promptId":"ae94540a-045a-4847-b16d-c2358b5c9db6","agentId":"a07e7799acd8be342","type":"user","message":{"role":"user","content":[{"tool_use_id":"toolu_01R7QDVTJSeWgcktJJ8axMcH","type":"tool_result","content":"23: c1_train: seq1_cluster array for training set\n24: c2_train: seq2_cluster array for training set\n25: c1_val: seq1_cluster array for validation set\n26: c2_val: seq2_cluster array for validation set\n36: # 2. Filter invalid clusters (-1)\n41: # 3. Validate cluster IDs don't exceed multiplier (prevent hash collisions)\n42: max_cluster = max(c1.max(), c2.max()) if len(c1) > 0 else 0\n43: assert max_cluster < multiplier, (\n44: f\"Cluster ID {max_cluster} exceeds multiplier {multiplier}. \"\n45: f\"Hash collisions will occur! Increase multiplier or reduce cluster IDs.\"\n55: print(\"WARNING: No valid cluster pairs found. All entries have cluster_id=-1.\")\n69: Pool 0 (PPI): cluster-weighted — sample a cluster uniformly, then an example from it.\n88: # Build cluster-weighted index for PPI examples (c1[i] >= 0 and c2[i] >= 0)\n89: cluster_to_indices = defaultdict(set)\n92: cluster_to_indices[int(c1[idx])].add(idx)\n93: cluster_to_indices[int(c2[idx])].add(idx)\n94: self.clusters = list(cluster_to_indices.keys())\n95: self.cluster_to_indices = {k: np.array(list(v), dtype=np.int64) for k, v in cluster_to_indices.items()}\n119: self.has_clusters = len(self.clusters) > 0\n138: # PPI: cluster-weighted sampling\n139: if self.has_clusters:\n140: cluster_idx = torch.randint(len(self.clusters), (1,), generator=g).item()\n141: cluster = self.clusters[cluster_idx]\n142: cluster_examples = self.cluster_to_indices[cluster]\n143: idx = cluster_examples[torch.randint(len(cluster_examples), (1,), generator=g).item()]\n162:def create_contact_targets(\n163: seq1_lens: List[int],\n164: seq2_lens: List[int],\n165: contacts_list: List[List],\n169: Create contact target tensor from list of contacts.\n172: seq1_lens: List of sequence 1 lengths\n173: seq2_lens: List of sequence 2 lengths\n174: contacts_list: List of contact lists, each contact is [i, j]\n178: contact_targets: (batch_size, target_map_size, target_map_size) bool tensor\n180: batch_size = len(seq1_lens)\n181: contact_targets = torch.zeros(batch_size, target_map_size, target_map_size, dtype=torch.bool)\n187: for i, (seq1_len, seq2_len, contacts) in enumerate(zip(seq1_lens, seq2_lens, contacts_list)):\n188: seq1_len = min(seq1_len, target_map_size)\n189: seq2_len = min(seq2_len, target_map_size)\n191: if contacts:\n192: contacts_tensor = torch.tensor(contacts, dtype=torch.long)\n193: valid_mask = (contacts_tensor[:, 0] < seq1_len) & (contacts_tensor[:, 1] < seq2_len)\n194: valid_contacts = contacts_tensor[valid_mask]\n196: if len(valid_contacts) > 0:\n197: batch_indices_list.append(torch.full((len(valid_contacts),), i, dtype=torch.long))\n198: r1_indices_list.append(valid_contacts[:, 0])\n199: r2_indices_list.append(valid_contacts[:, 1])\n205: contact_targets[batch_indices, r1_indices, r2_indices] = True\n207: return contact_targets\n213: - Tokenizes seq1 and seq2 separately.","is_error":false}]},"uuid":"09992d5e-6147-41f8-9039-2e02e15559d0","timestamp":"2026-06-01T15:28:39.432Z","sourceToolAssistantUUID":"e3920087-d844-46fa-b9cc-916f22123c21","userType":"external","entrypoint":"claude-vscode","cwd":"/home/ubuntu/ppi_project","sessionId":"5a87a0cd-6802-4ac5-a708-29b2bf7ea563","version":"2.1.159","gitBranch":"afdb"} +{"parentUuid":"09992d5e-6147-41f8-9039-2e02e15559d0","isSidechain":true,"agentId":"a07e7799acd8be342","message":{"model":"claude-haiku-4-5-20251001","id":"msg_017PqaPfhZbwoqNjVqYKUU3Y","type":"message","role":"assistant","content":[{"type":"text","text":"Perfect! Now let me check how samples are structured in the dataset more closely:"}],"stop_reason":null,"stop_sequence":null,"stop_details":null,"usage":{"input_tokens":4,"cache_creation_input_tokens":1147,"cache_read_input_tokens":66806,"cache_creation":{"ephemeral_5m_input_tokens":1147,"ephemeral_1h_input_tokens":0},"output_tokens":1,"service_tier":"standard","inference_geo":"not_available"},"diagnostics":null},"requestId":"req_011Cbcn1jo6E8aaduaqCxnu5","attributionAgent":"Explore","type":"assistant","uuid":"465a844c-d0fb-49eb-8c4d-6b927b281d6b","timestamp":"2026-06-01T15:28:41.149Z","userType":"external","entrypoint":"claude-vscode","cwd":"/home/ubuntu/ppi_project","sessionId":"5a87a0cd-6802-4ac5-a708-29b2bf7ea563","version":"2.1.159","gitBranch":"afdb"} +{"parentUuid":"465a844c-d0fb-49eb-8c4d-6b927b281d6b","isSidechain":true,"agentId":"a07e7799acd8be342","message":{"model":"claude-haiku-4-5-20251001","id":"msg_017PqaPfhZbwoqNjVqYKUU3Y","type":"message","role":"assistant","content":[{"type":"tool_use","id":"toolu_01JBrX4CF4HK2F6KSLQA3PPT","name":"Bash","input":{"command":"grep -A 30 \"class DataCollatorForPPI\" /home/ubuntu/ppi_project/dataset.py | head -50"},"caller":{"type":"direct"}}],"stop_reason":"tool_use","stop_sequence":null,"stop_details":null,"usage":{"input_tokens":4,"cache_creation_input_tokens":1147,"cache_read_input_tokens":66806,"output_tokens":100,"server_tool_use":{"web_search_requests":0,"web_fetch_requests":0},"service_tier":"standard","cache_creation":{"ephemeral_1h_input_tokens":0,"ephemeral_5m_input_tokens":1147},"inference_geo":"not_available","iterations":[{"input_tokens":4,"output_tokens":100,"cache_read_input_tokens":66806,"cache_creation_input_tokens":1147,"cache_creation":{"ephemeral_5m_input_tokens":1147,"ephemeral_1h_input_tokens":0},"type":"message"}],"speed":"standard"},"diagnostics":null},"requestId":"req_011Cbcn1jo6E8aaduaqCxnu5","attributionAgent":"Explore","type":"assistant","uuid":"bacb7da9-7a89-4b68-bada-9ce060a69dd8","timestamp":"2026-06-01T15:28:41.427Z","userType":"external","entrypoint":"claude-vscode","cwd":"/home/ubuntu/ppi_project","sessionId":"5a87a0cd-6802-4ac5-a708-29b2bf7ea563","version":"2.1.159","gitBranch":"afdb"} +{"parentUuid":"bacb7da9-7a89-4b68-bada-9ce060a69dd8","isSidechain":true,"promptId":"ae94540a-045a-4847-b16d-c2358b5c9db6","agentId":"a07e7799acd8be342","type":"user","message":{"role":"user","content":[{"tool_use_id":"toolu_01JBrX4CF4HK2F6KSLQA3PPT","type":"tool_result","content":"class DataCollatorForPPI:\n \"\"\"\n Custom data collator.\n - Tokenizes seq1 and seq2 separately.\n - Creates the 2D contact map target tensor.\n - Optionally swaps query/target with 50% probability.\n \"\"\"\n def __init__(self, tokenizer: AutoTokenizer, max_len: int, swap_augment: bool = True, adds_special_tokens: bool = True):\n self.tokenizer = tokenizer\n self.max_len = max_len\n self.swap_augment = swap_augment\n self.adds_special_tokens = adds_special_tokens # ESM adds CLS/SEP, gLM2 does not\n\n def __call__(self, batch: list[dict]):\n # Randomly swap query/target for augmentation\n if self.swap_augment:\n batch = [\n {'seq1': item['seq2'],\n 'seq2': item['seq1'], \n 'contacts': [(j, i) for i, j in item.get('contacts', [])],\n 'seq1_cluster': item.get('seq2_cluster'),\n 'seq2_cluster': item.get('seq1_cluster'),\n 'complex_id': item.get('complex_id')}\n if torch.rand(1).item() < 0.5 else item\n for item in batch\n ]\n \n # Extract sequences\n seqs1 = [item['seq1'] for item in batch]\n seqs2 = [item['seq2'] for item in batch]","is_error":false}]},"uuid":"3bc0d00d-f34a-4412-98b6-8351dbc99afe","timestamp":"2026-06-01T15:28:41.452Z","sourceToolAssistantUUID":"bacb7da9-7a89-4b68-bada-9ce060a69dd8","userType":"external","entrypoint":"claude-vscode","cwd":"/home/ubuntu/ppi_project","sessionId":"5a87a0cd-6802-4ac5-a708-29b2bf7ea563","version":"2.1.159","gitBranch":"afdb"} +{"parentUuid":"3bc0d00d-f34a-4412-98b6-8351dbc99afe","isSidechain":true,"agentId":"a07e7799acd8be342","message":{"model":"claude-haiku-4-5-20251001","id":"msg_01KqxEfausJ9JoXiE5wodnyv","type":"message","role":"assistant","content":[{"type":"text","text":"Let me get a comprehensive summary by searching for how data items are expected to be structured:"}],"stop_reason":null,"stop_sequence":null,"stop_details":null,"usage":{"input_tokens":4,"cache_creation_input_tokens":498,"cache_read_input_tokens":67953,"cache_creation":{"ephemeral_5m_input_tokens":498,"ephemeral_1h_input_tokens":0},"output_tokens":3,"service_tier":"standard","inference_geo":"not_available"},"diagnostics":null},"requestId":"req_011Cbcn1tS1H847FKRC1mCtq","attributionAgent":"Explore","type":"assistant","uuid":"861c604d-9971-428f-aca9-30d0660c98df","timestamp":"2026-06-01T15:28:42.968Z","userType":"external","entrypoint":"claude-vscode","cwd":"/home/ubuntu/ppi_project","sessionId":"5a87a0cd-6802-4ac5-a708-29b2bf7ea563","version":"2.1.159","gitBranch":"afdb"} +{"parentUuid":"861c604d-9971-428f-aca9-30d0660c98df","isSidechain":true,"agentId":"a07e7799acd8be342","message":{"model":"claude-haiku-4-5-20251001","id":"msg_01KqxEfausJ9JoXiE5wodnyv","type":"message","role":"assistant","content":[{"type":"tool_use","id":"toolu_01Ninwz24Pi2z7xS8CrC9bgC","name":"Bash","input":{"command":"grep -n \"item\\[\" /home/ubuntu/ppi_project/dataset.py | head -40"},"caller":{"type":"direct"}}],"stop_reason":null,"stop_sequence":null,"stop_details":null,"usage":{"input_tokens":4,"cache_creation_input_tokens":498,"cache_read_input_tokens":67953,"cache_creation":{"ephemeral_5m_input_tokens":498,"ephemeral_1h_input_tokens":0},"output_tokens":3,"service_tier":"standard","inference_geo":"not_available"},"diagnostics":null},"requestId":"req_011Cbcn1tS1H847FKRC1mCtq","attributionAgent":"Explore","type":"assistant","uuid":"1c642626-0db8-4410-93c4-0bc194551cfa","timestamp":"2026-06-01T15:28:43.014Z","userType":"external","entrypoint":"claude-vscode","cwd":"/home/ubuntu/ppi_project","sessionId":"5a87a0cd-6802-4ac5-a708-29b2bf7ea563","version":"2.1.159","gitBranch":"afdb"} +{"parentUuid":"1c642626-0db8-4410-93c4-0bc194551cfa","isSidechain":true,"promptId":"ae94540a-045a-4847-b16d-c2358b5c9db6","agentId":"a07e7799acd8be342","type":"user","message":{"role":"user","content":[{"tool_use_id":"toolu_01Ninwz24Pi2z7xS8CrC9bgC","type":"tool_result","content":"227: {'seq1': item['seq2'],\n228: 'seq2': item['seq1'], \n238: seqs1 = [item['seq1'] for item in batch]\n239: seqs2 = [item['seq2'] for item in batch]\n265: seq1_lens = [len(item['seq1']) for item in batch]\n266: seq2_lens = [len(item['seq2']) for item in batch]\n296: seqs1 = [item['seq1'] for item in batch]\n297: seqs2 = [item['seq2'] for item in batch]\n298: labels = torch.tensor([item['label'] for item in batch], dtype=torch.long)","is_error":false}]},"uuid":"55f1cb62-6acd-434e-8a95-0568c658045b","timestamp":"2026-06-01T15:28:43.045Z","sourceToolAssistantUUID":"1c642626-0db8-4410-93c4-0bc194551cfa","userType":"external","entrypoint":"claude-vscode","cwd":"/home/ubuntu/ppi_project","sessionId":"5a87a0cd-6802-4ac5-a708-29b2bf7ea563","version":"2.1.159","gitBranch":"afdb"} +{"parentUuid":"55f1cb62-6acd-434e-8a95-0568c658045b","isSidechain":true,"agentId":"a07e7799acd8be342","message":{"model":"claude-haiku-4-5-20251001","id":"msg_019JYkRAse6N3Lfhv4bqxG28","type":"message","role":"assistant","content":[{"type":"text","text":"Let me look at a sample of what the dataset collator returns:"}],"stop_reason":null,"stop_sequence":null,"stop_details":null,"usage":{"input_tokens":3,"cache_creation_input_tokens":283,"cache_read_input_tokens":68451,"cache_creation":{"ephemeral_5m_input_tokens":283,"ephemeral_1h_input_tokens":0},"output_tokens":5,"service_tier":"standard","inference_geo":"not_available"},"diagnostics":null},"requestId":"req_011Cbcn21Zsmihnn4bPVe5JK","attributionAgent":"Explore","type":"assistant","uuid":"87e808e9-229c-4f5d-b48d-98485aa1624a","timestamp":"2026-06-01T15:28:45.101Z","userType":"external","entrypoint":"claude-vscode","cwd":"/home/ubuntu/ppi_project","sessionId":"5a87a0cd-6802-4ac5-a708-29b2bf7ea563","version":"2.1.159","gitBranch":"afdb"} +{"parentUuid":"87e808e9-229c-4f5d-b48d-98485aa1624a","isSidechain":true,"agentId":"a07e7799acd8be342","message":{"model":"claude-haiku-4-5-20251001","id":"msg_019JYkRAse6N3Lfhv4bqxG28","type":"message","role":"assistant","content":[{"type":"tool_use","id":"toolu_01BWDHinJjA5NZSjE36ifnsU","name":"Read","input":{"file_path":"/home/ubuntu/ppi_project/dataset.py","offset":210,"limit":80},"caller":{"type":"direct"}}],"stop_reason":null,"stop_sequence":null,"stop_details":null,"usage":{"input_tokens":3,"cache_creation_input_tokens":283,"cache_read_input_tokens":68451,"cache_creation":{"ephemeral_5m_input_tokens":283,"ephemeral_1h_input_tokens":0},"output_tokens":5,"service_tier":"standard","inference_geo":"not_available"},"diagnostics":null},"requestId":"req_011Cbcn21Zsmihnn4bPVe5JK","attributionAgent":"Explore","type":"assistant","uuid":"1b44868c-bb35-48e2-b50b-2f0d990ab6e9","timestamp":"2026-06-01T15:28:45.256Z","userType":"external","entrypoint":"claude-vscode","cwd":"/home/ubuntu/ppi_project","sessionId":"5a87a0cd-6802-4ac5-a708-29b2bf7ea563","version":"2.1.159","gitBranch":"afdb"} +{"parentUuid":"1b44868c-bb35-48e2-b50b-2f0d990ab6e9","isSidechain":true,"promptId":"ae94540a-045a-4847-b16d-c2358b5c9db6","agentId":"a07e7799acd8be342","type":"user","message":{"role":"user","content":[{"tool_use_id":"toolu_01BWDHinJjA5NZSjE36ifnsU","type":"tool_result","content":"210\tclass DataCollatorForPPI:\n211\t \"\"\"\n212\t Custom data collator.\n213\t - Tokenizes seq1 and seq2 separately.\n214\t - Creates the 2D contact map target tensor.\n215\t - Optionally swaps query/target with 50% probability.\n216\t \"\"\"\n217\t def __init__(self, tokenizer: AutoTokenizer, max_len: int, swap_augment: bool = True, adds_special_tokens: bool = True):\n218\t self.tokenizer = tokenizer\n219\t self.max_len = max_len\n220\t self.swap_augment = swap_augment\n221\t self.adds_special_tokens = adds_special_tokens # ESM adds CLS/SEP, gLM2 does not\n222\t\n223\t def __call__(self, batch: list[dict]):\n224\t # Randomly swap query/target for augmentation\n225\t if self.swap_augment:\n226\t batch = [\n227\t {'seq1': item['seq2'],\n228\t 'seq2': item['seq1'], \n229\t 'contacts': [(j, i) for i, j in item.get('contacts', [])],\n230\t 'seq1_cluster': item.get('seq2_cluster'),\n231\t 'seq2_cluster': item.get('seq1_cluster'),\n232\t 'complex_id': item.get('complex_id')}\n233\t if torch.rand(1).item() < 0.5 else item\n234\t for item in batch\n235\t ]\n236\t \n237\t # Extract sequences\n238\t seqs1 = [item['seq1'] for item in batch]\n239\t seqs2 = [item['seq2'] for item in batch]\n240\t \n241\t # Normalize cluster IDs -1 (use -1 for missing/DDI dataset)\n242\t seq1_clusters = [item.get('seq1_cluster') for item in batch]\n243\t seq1_clusters = [-1 if val is None else val for val in seq1_clusters]\n244\t seq2_clusters = [item.get('seq2_cluster') for item in batch]\n245\t seq2_clusters = [-1 if val is None else val for val in seq2_clusters]\n246\t \n247\t # Tokenize separately (ESM adds CLS/SEP, gLM2 does not)\n248\t batch1_tokens = self.tokenizer(\n249\t seqs1, \n250\t return_tensors='pt', \n251\t padding='max_length', \n252\t truncation=True, \n253\t max_length=self.max_len\n254\t )\n255\t \n256\t batch2_tokens = self.tokenizer(\n257\t seqs2, \n258\t return_tensors='pt', \n259\t padding='max_length', \n260\t truncation=True, \n261\t max_length=self.max_len\n262\t )\n263\t \n264\t # Create contact targets using standalone function\n265\t seq1_lens = [len(item['seq1']) for item in batch]\n266\t seq2_lens = [len(item['seq2']) for item in batch]\n267\t # The target map size matches model output (which slices off CLS/SEP if present)\n268\t target_map_size = self.max_len - 2 if self.adds_special_tokens else self.max_len\n269\t contacts_list = [item.get('contacts', []) for item in batch]\n270\t contact_targets = create_contact_targets(seq1_lens, seq2_lens, contacts_list, target_map_size)\n271\t\n272\t # Carry is_positive flag if present (defaults to True for backward compatibility/PDB_PPI)\n273\t is_positive = [item.get('is_positive', True) for item in batch]\n274\t\n275\t return {\n276\t 'seq1_tokens': batch1_tokens,\n277\t 'seq2_tokens': batch2_tokens,\n278\t 'contact_target': contact_targets,\n279\t 'seq1_cluster': torch.tensor(seq1_clusters, dtype=torch.long),\n280\t 'seq2_cluster': torch.tensor(seq2_clusters, dtype=torch.long),\n281\t 'is_positive': torch.tensor(is_positive, dtype=torch.bool)\n282\t }\n283\t\n284\t\n285\tclass BenchmarkCollator:\n286\t \"\"\"\n287\t Simple data collator for benchmark evaluation.\n288\t Tokenizes seq1 and seq2 separately and includes labels.\n289\t \"\"\""}]},"uuid":"4728c11b-f6ec-4bb2-bb38-c677c6fd082c","timestamp":"2026-06-01T15:28:45.263Z","sourceToolAssistantUUID":"1b44868c-bb35-48e2-b50b-2f0d990ab6e9","userType":"external","entrypoint":"claude-vscode","cwd":"/home/ubuntu/ppi_project","sessionId":"5a87a0cd-6802-4ac5-a708-29b2bf7ea563","version":"2.1.159","gitBranch":"afdb"} +{"parentUuid":"4728c11b-f6ec-4bb2-bb38-c677c6fd082c","isSidechain":true,"agentId":"a07e7799acd8be342","message":{"model":"claude-haiku-4-5-20251001","id":"msg_01H1kMpMxX1U9597P1SwnuM6","type":"message","role":"assistant","content":[{"type":"text","text":"Perfect! Now let me check the model forward pass to understand how data flows through:"}],"stop_reason":null,"stop_sequence":null,"stop_details":null,"usage":{"input_tokens":3,"cache_creation_input_tokens":1299,"cache_read_input_tokens":68734,"cache_creation":{"ephemeral_5m_input_tokens":1299,"ephemeral_1h_input_tokens":0},"output_tokens":1,"service_tier":"standard","inference_geo":"not_available"},"diagnostics":null},"requestId":"req_011Cbcn2Bwxa5ZjtyNrgBGrM","attributionAgent":"Explore","type":"assistant","uuid":"3fd43cc8-f400-4269-b48b-765aac2403e3","timestamp":"2026-06-01T15:28:47.950Z","userType":"external","entrypoint":"claude-vscode","cwd":"/home/ubuntu/ppi_project","sessionId":"5a87a0cd-6802-4ac5-a708-29b2bf7ea563","version":"2.1.159","gitBranch":"afdb"} +{"parentUuid":"3fd43cc8-f400-4269-b48b-765aac2403e3","isSidechain":true,"agentId":"a07e7799acd8be342","message":{"model":"claude-haiku-4-5-20251001","id":"msg_01H1kMpMxX1U9597P1SwnuM6","type":"message","role":"assistant","content":[{"type":"tool_use","id":"toolu_01UAzztB67s878gyvahiPnpk","name":"Bash","input":{"command":"grep -n \"batch\\[\" /home/ubuntu/ppi_project/train.py | head -30"},"caller":{"type":"direct"}}],"stop_reason":null,"stop_sequence":null,"stop_details":null,"usage":{"input_tokens":3,"cache_creation_input_tokens":1299,"cache_read_input_tokens":68734,"cache_creation":{"ephemeral_5m_input_tokens":1299,"ephemeral_1h_input_tokens":0},"output_tokens":1,"service_tier":"standard","inference_geo":"not_available"},"diagnostics":null},"requestId":"req_011Cbcn2Bwxa5ZjtyNrgBGrM","attributionAgent":"Explore","type":"assistant","uuid":"078b07c9-40cc-402f-a174-b3ce4407952e","timestamp":"2026-06-01T15:28:48.560Z","userType":"external","entrypoint":"claude-vscode","cwd":"/home/ubuntu/ppi_project","sessionId":"5a87a0cd-6802-4ac5-a708-29b2bf7ea563","version":"2.1.159","gitBranch":"afdb"} +{"parentUuid":"078b07c9-40cc-402f-a174-b3ce4407952e","isSidechain":true,"promptId":"ae94540a-045a-4847-b16d-c2358b5c9db6","agentId":"a07e7799acd8be342","type":"user","message":{"role":"user","content":[{"tool_use_id":"toolu_01UAzztB67s878gyvahiPnpk","type":"tool_result","content":"495: seq1_ids = batch['seq1_tokens']['input_ids']\n496: seq1_mask = batch['seq1_tokens']['attention_mask']\n497: seq2_ids = batch['seq2_tokens']['input_ids']\n498: seq2_mask = batch['seq2_tokens']['attention_mask']\n499: contact_target = batch['contact_target']","is_error":false}]},"uuid":"61132293-d567-4295-a1c3-427d2f6e7ce0","timestamp":"2026-06-01T15:28:48.587Z","sourceToolAssistantUUID":"078b07c9-40cc-402f-a174-b3ce4407952e","userType":"external","entrypoint":"claude-vscode","cwd":"/home/ubuntu/ppi_project","sessionId":"5a87a0cd-6802-4ac5-a708-29b2bf7ea563","version":"2.1.159","gitBranch":"afdb"} +{"parentUuid":"61132293-d567-4295-a1c3-427d2f6e7ce0","isSidechain":true,"agentId":"a07e7799acd8be342","message":{"model":"claude-haiku-4-5-20251001","id":"msg_0183o6eQnXdGEFYGi5eCvTWx","type":"message","role":"assistant","content":[{"type":"text","text":"Excellent! Now let me create a comprehensive summary. Let me just verify the exact shapes and formats by looking at one more key section:"}],"stop_reason":null,"stop_sequence":null,"stop_details":null,"usage":{"input_tokens":3,"cache_creation_input_tokens":199,"cache_read_input_tokens":70033,"cache_creation":{"ephemeral_5m_input_tokens":199,"ephemeral_1h_input_tokens":0},"output_tokens":7,"service_tier":"standard","inference_geo":"not_available"},"diagnostics":null},"requestId":"req_011Cbcn2RX2U66NCrp1RJENp","attributionAgent":"Explore","type":"assistant","uuid":"9a3ac791-b4c5-469e-80dd-c287ac1cf000","timestamp":"2026-06-01T15:28:51.154Z","userType":"external","entrypoint":"claude-vscode","cwd":"/home/ubuntu/ppi_project","sessionId":"5a87a0cd-6802-4ac5-a708-29b2bf7ea563","version":"2.1.159","gitBranch":"afdb"} +{"parentUuid":"9a3ac791-b4c5-469e-80dd-c287ac1cf000","isSidechain":true,"agentId":"a07e7799acd8be342","message":{"model":"claude-haiku-4-5-20251001","id":"msg_0183o6eQnXdGEFYGi5eCvTWx","type":"message","role":"assistant","content":[{"type":"tool_use","id":"toolu_01EAWvNFXjPg2fUw1CtxiMg9","name":"Read","input":{"file_path":"/home/ubuntu/ppi_project/model.py","offset":357,"limit":40},"caller":{"type":"direct"}}],"stop_reason":"tool_use","stop_sequence":null,"stop_details":null,"usage":{"input_tokens":3,"cache_creation_input_tokens":199,"cache_read_input_tokens":70033,"output_tokens":128,"server_tool_use":{"web_search_requests":0,"web_fetch_requests":0},"service_tier":"standard","cache_creation":{"ephemeral_1h_input_tokens":0,"ephemeral_5m_input_tokens":199},"inference_geo":"not_available","iterations":[{"input_tokens":3,"output_tokens":128,"cache_read_input_tokens":70033,"cache_creation_input_tokens":199,"cache_creation":{"ephemeral_5m_input_tokens":199,"ephemeral_1h_input_tokens":0},"type":"message"}],"speed":"standard"},"diagnostics":null},"requestId":"req_011Cbcn2RX2U66NCrp1RJENp","attributionAgent":"Explore","type":"assistant","uuid":"3a3e29db-e58e-4b4b-8597-4bd03867f464","timestamp":"2026-06-01T15:28:51.588Z","userType":"external","entrypoint":"claude-vscode","cwd":"/home/ubuntu/ppi_project","sessionId":"5a87a0cd-6802-4ac5-a708-29b2bf7ea563","version":"2.1.159","gitBranch":"afdb"} +{"parentUuid":"3a3e29db-e58e-4b4b-8597-4bd03867f464","isSidechain":true,"promptId":"ae94540a-045a-4847-b16d-c2358b5c9db6","agentId":"a07e7799acd8be342","type":"user","message":{"role":"user","content":[{"tool_use_id":"toolu_01EAWvNFXjPg2fUw1CtxiMg9","type":"tool_result","content":"357\t\n358\t def encode_protein(\n359\t self,\n360\t input_ids: torch.Tensor,\n361\t attention_mask: Optional[torch.Tensor] = None,\n362\t ):\n363\t \"\"\"\n364\t Runs the shared plm model on tokenized protein sequences.\n365\t \n366\t Args:\n367\t input_ids (torch.Tensor): (B, L) - Token IDs from the tokenizer\n368\t attention_mask (Optional[torch.Tensor]): (B, L) - Attention mask (1 for real tokens, 0 for padding)\n369\t \n370\t Returns:\n371\t residue_embeds (torch.Tensor): (B, L, D_plm) - Embeddings for each residue\n372\t pool_embed (torch.Tensor): (B, D_plm) - Pooled embedding using mask-aware pooling\n373\t \"\"\"\n374\t outputs = self.plm(input_ids=input_ids, attention_mask=attention_mask)\n375\t \n376\t # (B, L, D_plm) - Full sequence embeddings\n377\t residue_embeds = outputs.last_hidden_state\n378\t return residue_embeds\n379\t\n380\t def predict_contacts(self, embed1: torch.Tensor, embed2: torch.Tensor, mask1: torch.Tensor, mask2: torch.Tensor):\n381\t \"\"\"\n382\t Predicts a contact map using dot product between projected residue embeddings.\n383\t \n384\t Args:\n385\t embed1 (torch.Tensor): (B, L1, D_plm) - Residue embeddings for sequence 1 (Query)\n386\t embed2 (torch.Tensor): (B, L2, D_plm) - Residue embeddings for sequence 2 (Key)\n387\t mask1 (torch.Tensor): (B, L1) - Attention mask for seq1\n388\t mask2 (torch.Tensor): (B, L2) - Attention mask for seq2\n389\t \n390\t Returns:\n391\t contact_logits (torch.Tensor): (B, L1, L2) - Contact prediction logits\n392\t \"\"\"\n393\t return self.contact_head(embed1, embed2, mask1, mask2)\n394\t \n395\t def encode_pair(\n396\t self,"}]},"uuid":"ad3ae9e7-da3f-4b7a-a409-141410e447cf","timestamp":"2026-06-01T15:28:51.593Z","sourceToolAssistantUUID":"3a3e29db-e58e-4b4b-8597-4bd03867f464","userType":"external","entrypoint":"claude-vscode","cwd":"/home/ubuntu/ppi_project","sessionId":"5a87a0cd-6802-4ac5-a708-29b2bf7ea563","version":"2.1.159","gitBranch":"afdb"} +{"parentUuid":"ad3ae9e7-da3f-4b7a-a409-141410e447cf","isSidechain":true,"agentId":"a07e7799acd8be342","message":{"model":"claude-haiku-4-5-20251001","id":"msg_01W3xBZYuDE7PPP7ivpbS2kq","type":"message","role":"assistant","content":[{"type":"text","text":"Perfect! Now I have comprehensive information. Let me create a final detailed summary:\n\n## Summary: Data Loading and Training Pipeline\n\n### 1. **Dataset/Dataloader Files and Sample Fields**\n\n**Primary files:**\n- `/home/ubuntu/ppi_project/dataset.py` - Contains dataset classes, collators, and sampling logic\n- `/home/ubuntu/ppi_project/train.py` - Main training script\n- `/home/ubuntu/ppi_project/model.py` - Model architecture\n\n**Sample structure (from HuggingFace datasets):**\nEach sample in the dataset is a dictionary with the following fields:\n- `seq1` (str) - First protein sequence (amino acid string)\n- `seq2` (str) - Second protein sequence (amino acid string)\n- `seq1_cluster` (int) - Cluster ID for sequence 1 (or -1 if not available, e.g., DDI dataset)\n- `seq2_cluster` (int) - Cluster ID for sequence 2 (or -1 if not available)\n- `contacts` (list of [int, int]) - List of contact pairs as residue indices: `[[i1, j1], [i2, j2], ...]` where i ∈ [0, len(seq1)-1], j ∈ [0, len(seq2)-1]\n- `complex_id` (str, optional) - Complex identifier\n- `is_positive` (bool, added by train.py) - Flag indicating if this is a positive/known interaction\n\n### 2. **Sequence Tokenization and Padding**\n\n**Location:** `/home/ubuntu/ppi_project/dataset.py`, lines 248-262 (in `DataCollatorForPPI.__call__`)\n\n**Process:**\n- Sequences are tokenized **separately** for seq1 and seq2\n- Uses `transformers.AutoTokenizer` from the PLM model name (specified in config as `plm_model_name`)\n- Both sequences are padded/truncated to `max_len` (configurable, default 512)\n- Tokenization returns `input_ids` and `attention_mask`\n\n**Key detail:** \n- ESM models add CLS/SEP special tokens automatically\n- gLM2 models do not add special tokens\n- The config tracks this with `adds_special_tokens` flag (set based on model name in `train.py`, line 235)\n\n**Returned tokenized data structure:**\n```python\n'seq1_tokens': {\n 'input_ids': (B, max_len), # Token IDs\n 'attention_mask': (B, max_len) # 1 for real tokens, 0 for padding\n}\n'seq2_tokens': {\n 'input_ids': (B, max_len),\n 'attention_mask': (B, max_len)\n}\n```\n\n### 3. **Contact Map Format and Use**\n\n**Location:** `/home/ubuntu/ppi_project/dataset.py`, lines 162-207 (function `create_contact_targets`)\n\n**Format:**\n- Contact maps are created as **3D boolean tensors**: `(batch_size, target_map_size, target_map_size)`\n- `target_map_size = max_len - 2` if ESM (CLS/SEP removed), else `max_len`\n- Shape at model output: `(B, L1-adjusted, L2-adjusted)` where L1, L2 are within target_map_size\n\n**Creation process:**\n1. Raw contacts come as list of `[i, j]` indices from dataset samples\n2. For each pair in a batch:\n - Valid contacts are those where `contact[0] < seq1_len` and `contact[1] < seq2_len`\n - Valid contacts are set to `True` in the 3D tensor\n - Padding positions remain `False`\n3. Result: Boolean tensor where `contact_target[b, i, j] = True` if residue i in seq1 contacts residue j in seq2 in batch element b\n\n**Usage in loss:**\n- `/home/ubuntu/ppi_project/loss.py`, `ContactLoss.forward()` (lines 89-117)\n- Contacts are compared against `contact_logits` from the model\n- Loss is computed only on valid (non-padded) positions using `valid_mask`\n\n### 4. **Training Loop and Collate Function**\n\n**Collate function:** `DataCollatorForPPI` (lines 210-282 in `dataset.py`)\n\n**Features:**\n- **Swap augmentation (50% probability):** Query/target sequences are randomly swapped, with contacts transposed: `(i, j) → (j, i)`\n- **Cluster ID normalization:** Missing cluster IDs become -1\n- **Multi-dataset support:** Handles PPI, DDI, AFDB, and negative datasets with different cluster ID schemes\n- **Positive flag tracking:** Carries `is_positive` to distinguish true interactions from negatives\n\n**Training loop:** `/home/ubuntu/ppi_project/train.py`, lines 481-622\n\n**Key components:**\n1. **Data loading** (lines 119-220): Loads train/val splits from HF datasets, concatenates optional DDI/AFDB datasets\n2. **Sampling** (lines 243-274): \n - Optional cluster-weighted sampler (`ClusterSampler`) for balanced sampling from PPI clusters\n - DDI and AFDB pools can be sampled with explicit probabilities\n3. **Batch processing** (lines 488-562):\n - Extracts tokenized inputs: `batch['seq1_tokens']['input_ids']`, etc.\n - Passes to model forward\n - Computes CLIP loss (retrieval) and contact prediction loss\n - **With GradCache** (line 489): Uses gradient accumulation for larger effective batches\n - **Without GradCache** (line 494): Standard backprop with two separate loss components\n\n**Batch structure passed to model:**\n```python\nbatch = {\n 'seq1_tokens': {'input_ids': (B, max_len), 'attention_mask': (B, max_len)},\n 'seq2_tokens': {'input_ids': (B, max_len), 'attention_mask': (B, max_len)},\n 'contact_target': (B, target_map_size, target_map_size, dtype=bool),\n 'seq1_cluster': (B,) tensor of int64,\n 'seq2_cluster': (B,) tensor of int64,\n 'is_positive': (B,) tensor of bool\n}\n```\n\n### 5. **Cropping and Windowing Logic**\n\n**Location:** `/home/ubuntu/ppi_project/model.py`, lines 512-517\n\n**Special token slicing:**\nWhen ESM tokenizer is used (which adds CLS/SEP):\n```python\nif self._adds_special_tokens:\n residue_embed1 = residue_embed1[:, 1:-1, :] # Remove CLS and SEP\n residue_embed2 = residue_embed2[:, 1:-1, :]\n seq1_mask = seq1_mask[:, 1:-1] # Slice mask accordingly\n seq2_mask = seq2_mask[:, 1:-1]\n```\n\n**Contact map size adjustment:**\n- Tokenized sequences are padded to `max_len` (e.g., 512)\n- Contact target maps are created with `target_map_size = max_len - 2` for ESM (to account for CLS/SEP removal)\n- Contact target maps use `target_map_size = max_len` for gLM2 (no special tokens)\n- This ensures shapes align: model outputs (B, L1-actual, L2-actual) match contact targets (B, target_map_size, target_map_size)\n\n**No explicit sliding window:** The pipeline uses full-length sequences up to `max_len`, not windowing. Very long sequences are truncated.\n\n### 6. **Data Flow Summary**\n\n```\nRaw Sample (seq1, seq2, contacts, clusters)\n ↓\nDataCollatorForPPI (batch of samples)\n ├─→ Tokenize: seq1, seq2 → input_ids, attention_mask (both padded to max_len)\n ├─→ Create contact targets: raw contacts → (B, target_map_size, target_map_size) bool tensor\n └─→ Return: {seq1_tokens, seq2_tokens, contact_target, clusters, is_positive}\n ↓\nModel.forward(seq1_ids, seq2_ids, seq1_mask, seq2_mask, seq1_cluster, seq2_cluster)\n ├─→ PLM encoding: (B, max_len) → (B, max_len, D_plm)\n ├─→ Slice CLS/SEP if ESM: (B, max_len-2, D_plm)\n ├─→ Contact head: (B, L1, D_plm) × (B, L2, D_plm) → (B, L1, L2) logits\n └─→ Return: contact_logits, valid_mask, clip_embeddings, logit_scale\n ↓\nLoss computation:\n ├─→ ContactLoss(logits, targets, valid_mask) → scalar loss\n └─→ CLIPLoss(embeddings, cluster_ids) → scalar loss\n```\n\n### Key File Paths\n- `/home/ubuntu/ppi_project/dataset.py` - Sampler, collators, dataset utilities\n- `/home/ubuntu/ppi_project/train.py` - Training loop, data loading, batch processing (lines 488-562 for forward pass)\n- `/home/ubuntu/ppi_project/model.py` - `ContrastivePPIModel`, contact head architecture, forward pass (lines 466-561)\n- `/home/ubuntu/ppi_project/config.py` - `ModelConfig` and `TrainingConfig` dataclasses\n- `/home/ubuntu/ppi_project/loss.py` - Contact loss and CLIP loss implementations"}],"stop_reason":"end_turn","stop_sequence":null,"stop_details":null,"usage":{"input_tokens":3,"cache_creation_input_tokens":741,"cache_read_input_tokens":70232,"output_tokens":2446,"server_tool_use":{"web_search_requests":0,"web_fetch_requests":0},"service_tier":"standard","cache_creation":{"ephemeral_1h_input_tokens":0,"ephemeral_5m_input_tokens":741},"inference_geo":"not_available","iterations":[{"input_tokens":3,"output_tokens":2446,"cache_read_input_tokens":70232,"cache_creation_input_tokens":741,"cache_creation":{"ephemeral_5m_input_tokens":741,"ephemeral_1h_input_tokens":0},"type":"message"}],"speed":"standard"},"diagnostics":null},"requestId":"req_011Cbcn2dqhJoWpSa1pYcrMy","attributionAgent":"Explore","type":"assistant","uuid":"8a615a4d-9742-4e93-a776-a85a1b8c98af","timestamp":"2026-06-01T15:29:13.129Z","userType":"external","entrypoint":"claude-vscode","cwd":"/home/ubuntu/ppi_project","sessionId":"5a87a0cd-6802-4ac5-a708-29b2bf7ea563","version":"2.1.159","gitBranch":"afdb"} diff --git a/config.json b/config.json new file mode 100644 index 0000000000000000000000000000000000000000..bf84b946039e6e4162daef3f0add86721e2fe2a5 --- /dev/null +++ b/config.json @@ -0,0 +1,25 @@ +{ + "architectures": [ + "FlashPPIModel" + ], + "clip_embed_dim": 1024, + "contact_embed_dim": 1280, + "contact_num_heads": 8, + "contact_transformer_depth": 2, + "dtype": "float32", + "max_position_embeddings": 1024, + "model_type": "flashppi", + "plm_depth": 33, + "plm_dim": 1280, + "plm_ffn_dim_multiplier": null, + "plm_heads": 20, + "plm_norm_eps": 1e-05, + "plm_swiglu_multiple_of": 256, + "plm_vocab_size": 37, + "transformers_version": "4.57.3", + "use_flash_attention": true, + "auto_map": { + "AutoConfig": "configuration_flashppi.FlashPPIConfig", + "AutoModel": "modeling_flashppi.FlashPPIModel" + } +} \ No newline at end of file diff --git a/configuration_flashppi.py b/configuration_flashppi.py new file mode 100644 index 0000000000000000000000000000000000000000..622a04946d4290e1c3498d1c06462f39517cc781 --- /dev/null +++ b/configuration_flashppi.py @@ -0,0 +1,44 @@ +"""FlashPPI model configuration.""" + +from transformers import PretrainedConfig + + +class FlashPPIConfig(PretrainedConfig): + + model_type = "flashppi" + + def __init__( + self, + # gLM2 backbone config (defaults match gLM2_650M) + plm_dim: int = 1280, + plm_depth: int = 33, + plm_heads: int = 20, + plm_vocab_size: int = 37, + plm_norm_eps: float = 1e-5, + plm_swiglu_multiple_of: int = 256, + plm_ffn_dim_multiplier: float = None, + # FlashPPI head config + clip_embed_dim: int = 1024, + contact_embed_dim: int = 1280, + contact_num_heads: int = 8, + contact_transformer_depth: int = 2, + max_position_embeddings: int = 512, + use_flash_attention: bool = True, + **kwargs + ): + super().__init__(**kwargs) + # gLM2 config + self.plm_dim = plm_dim + self.plm_depth = plm_depth + self.plm_heads = plm_heads + self.plm_vocab_size = plm_vocab_size + self.plm_norm_eps = plm_norm_eps + self.plm_swiglu_multiple_of = plm_swiglu_multiple_of + self.plm_ffn_dim_multiplier = plm_ffn_dim_multiplier + # FlashPPI config + self.clip_embed_dim = clip_embed_dim + self.contact_embed_dim = contact_embed_dim + self.contact_num_heads = contact_num_heads + self.contact_transformer_depth = contact_transformer_depth + self.max_position_embeddings = max_position_embeddings + self.use_flash_attention = use_flash_attention diff --git a/data.mdb b/data.mdb new file mode 100644 index 0000000000000000000000000000000000000000..8ae7d0272fce6c549ee73fa9c93497713c112e22 Binary files /dev/null and b/data.mdb differ diff --git a/glm_tokenizer.py b/glm_tokenizer.py new file mode 100644 index 0000000000000000000000000000000000000000..ea39726a17e3637b49c86d6d71c2e1ca7d8df8dd --- /dev/null +++ b/glm_tokenizer.py @@ -0,0 +1,48 @@ +from tokenizers import Tokenizer +from tokenizers.models import BPE +from transformers import PreTrainedTokenizerFast + + +class gLM2Tokenizer(PreTrainedTokenizerFast): + + VOCAB = [ + "", "", "", "", + "L", "A", "G", "V", "S", "E", "R", "T", "I", "D", "P", "K", + "Q", "N", "F", "Y", "M", "H", "W", "C", "X", "B", "U", "Z", + "O", "a", "t", "c", "g", "<+>", "<->", "", "", + ] + + def __init__( + self, + unk_token="", + cls_token="", + pad_token="", + mask_token="", + eos_token="", + sep_token="", + pos_token="<+>", + neg_token="<->", + **kwargs, + ): + all_tokens = self.VOCAB + token_to_id = {tok: ind for ind, tok in enumerate(all_tokens)} + + bpe = BPE(token_to_id, merges=[], unk_token=str(unk_token)) + tokenizer = Tokenizer(bpe) + special_tokens = [cls_token, pad_token, + mask_token, eos_token, sep_token, pos_token, neg_token] + + tokenizer.add_special_tokens( + special_tokens, + ) + + super().__init__( + tokenizer_object=tokenizer, + unk_token=unk_token, + cls_token=cls_token, + pad_token=pad_token, + mask_token=mask_token, + eos_token=eos_token, + sep_token=sep_token, + **kwargs, + ) diff --git a/lock.mdb b/lock.mdb new file mode 100644 index 0000000000000000000000000000000000000000..cd82c6747623c71ae7ac587ec1d6ca692895cf79 Binary files /dev/null and b/lock.mdb differ diff --git a/model.safetensors b/model.safetensors new file mode 100644 index 0000000000000000000000000000000000000000..02399a2f75a2f9e9b61aa1edd05a9d302344639a --- /dev/null +++ b/model.safetensors @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:2d730369e54cd83c634aaa8d3b6d0adaf23ed8034483c7f1e2b242b516b94ac1 +size 2931379208 diff --git a/modeling_flashppi.py b/modeling_flashppi.py new file mode 100644 index 0000000000000000000000000000000000000000..78a557ce771a1ff7a08017c3aac2d919660daac1 --- /dev/null +++ b/modeling_flashppi.py @@ -0,0 +1,557 @@ +import torch +import torch.nn as nn +import torch.nn.functional as F +from dataclasses import dataclass +from typing import Optional, Tuple, Union +from einops import rearrange, repeat +from torch.utils.checkpoint import checkpoint +from transformers import PreTrainedModel +from transformers.modeling_outputs import ModelOutput + +from .configuration_flashppi import FlashPPIConfig + +# Detect Flash Attention installation +try: + from flash_attn.layers.rotary import apply_rotary_emb_func + from flash_attn import flash_attn_varlen_kvpacked_func + from flash_attn.bert_padding import pad_input, unpad_input + FLASH_ATTN_AVAILABLE = True +except ImportError: + FLASH_ATTN_AVAILABLE = False + unpad_input = pad_input = apply_rotary_emb_func = None + flash_attn_varlen_kvpacked_func = None + +def swiglu(x, y): + return F.silu(x) * y + +class RMSNorm(nn.Module): + """RMSNorm without variance_epsilon buffer for checkpoint compatibility.""" + def __init__(self, dim, eps=1e-6): + super().__init__() + self.weight = nn.Parameter(torch.ones(dim)) + self.eps = eps + + def forward(self, hidden_states): + input_dtype = hidden_states.dtype + hidden_states = hidden_states.to(torch.float32) + variance = hidden_states.pow(2).mean(-1, keepdim=True) + hidden_states = hidden_states * torch.rsqrt(variance + self.eps) + return (self.weight * hidden_states).to(input_dtype) + +@dataclass +class FlashPPIOutput(ModelOutput): + """Output type for FlashPPI model. + + Args: + contact_map: (B, L1, L2) contact probabilities between residue pairs. + contact_score: (B,) maximum contact probability per pair. + clip_embed1: (B, D) CLIP embedding for first protein. + clip_embed2: (B, D) CLIP embedding for second protein. + clip_score: (B,) CLIP similarity score (cosine similarity). + """ + contact_map: Optional[torch.FloatTensor] = None + contact_score: Optional[torch.FloatTensor] = None + clip_embed1: Optional[torch.FloatTensor] = None + clip_embed2: Optional[torch.FloatTensor] = None + clip_score: Optional[torch.FloatTensor] = None + + +def rotate_half(x, interleaved=False): + if not interleaved: + x1, x2 = x.chunk(2, dim=-1) + return torch.cat((-x2, x1), dim=-1) + else: + x1, x2 = x[..., ::2], x[..., 1::2] + return rearrange(torch.stack((-x2, x1), dim=-1), "... d two -> ... (d two)", two=2) + + +def apply_rotary_emb_torch(x, cos, sin, interleaved=False, position_ids=None): + """Apply rotary embeddings using pure PyTorch.""" + if position_ids is not None: + cos = cos[position_ids] + sin = sin[position_ids] + else: + cos = cos[:x.shape[1]] + sin = sin[:x.shape[1]] + + if not interleaved: + cos = repeat(cos, "... d -> ... 1 (2 d)") + sin = repeat(sin, "... d -> ... 1 (2 d)") + else: + cos = repeat(cos, "... d -> ... 1 (d 2)") + sin = repeat(sin, "... d -> ... 1 (d 2)") + + ro_dim = cos.shape[-1] + return torch.cat([ + x[..., :ro_dim] * cos + rotate_half(x[..., :ro_dim], interleaved) * sin, + x[..., ro_dim:], + ], dim=-1) + + +class RotaryEmbedding(nn.Module): + """Rotary position embeddings with flash attention support.""" + + def __init__(self, dim: int, base: float = 10000.0, interleaved: bool = False, device=None): + super().__init__() + self.dim = dim + self.base = float(base) + self.interleaved = interleaved + inv_freq = 1.0 / (self.base ** (torch.arange(0, dim, 2, device=device, dtype=torch.float32) / dim)) + self.register_buffer("inv_freq", inv_freq, persistent=False) + self._seq_len_cached = 0 + self._cos_cached = None + self._sin_cached = None + + def _update_cos_sin_cache(self, seqlen, device=None, dtype=None): + if seqlen > self._seq_len_cached or self._cos_cached is None or self._cos_cached.device != device: + self._seq_len_cached = seqlen + t = torch.arange(seqlen, device=device, dtype=torch.float32) + freqs = torch.outer(t, self.inv_freq.to(device=device, dtype=torch.float32)) + self._cos_cached = torch.cos(freqs).to(dtype) + self._sin_cached = torch.sin(freqs).to(dtype) + + def forward( + self, + q: torch.Tensor, + k: torch.Tensor, + cu_seqlens: Optional[torch.Tensor] = None, + max_seqlen: Optional[int] = None, + position_ids: Optional[torch.Tensor] = None, + ) -> Tuple[torch.Tensor, torch.Tensor]: + seqlen = q.shape[1] if max_seqlen is None else max_seqlen + self._update_cos_sin_cache(seqlen, device=q.device, dtype=q.dtype) + + if FLASH_ATTN_AVAILABLE and cu_seqlens is not None: + q = apply_rotary_emb_func( + q, self._cos_cached, self._sin_cached, + interleaved=self.interleaved, inplace=True, + cu_seqlens=cu_seqlens, max_seqlen=max_seqlen, + ) + k = apply_rotary_emb_func( + k, self._cos_cached, self._sin_cached, + interleaved=self.interleaved, inplace=True, + cu_seqlens=cu_seqlens, max_seqlen=max_seqlen, + ) + else: + q = apply_rotary_emb_torch(q, self._cos_cached, self._sin_cached, self.interleaved, position_ids) + k = apply_rotary_emb_torch(k, self._cos_cached, self._sin_cached, self.interleaved, position_ids) + return q, k + + +class Attention(nn.Module): + """Multi-head attention with optional flash attention.""" + + def __init__(self, dim: int, num_heads: int, use_rope: bool = True): + super().__init__() + self.n_heads = num_heads + self.head_dim = dim // num_heads + self.wqkv = nn.Linear(dim, num_heads * self.head_dim * 3, bias=False) + self.wo = nn.Linear(num_heads * self.head_dim, dim, bias=False) + self.rotary_emb = RotaryEmbedding(self.head_dim) if use_rope else None + + def forward( + self, + x: torch.Tensor, + cu_seqlens: Optional[torch.Tensor] = None, + max_seq_len: Optional[int] = None, + attention_mask: Optional[torch.Tensor] = None, + position_ids: Optional[torch.Tensor] = None, + ) -> torch.Tensor: + qkv = self.wqkv(x) + + if cu_seqlens is not None and FLASH_ATTN_AVAILABLE: + # Flash attention path (unpadded) + total_seqlen = x.shape[0] + q, k, v = torch.split(qkv, self.n_heads * self.head_dim, dim=-1) + q = q.view(total_seqlen, self.n_heads, self.head_dim) + k = k.view(total_seqlen, self.n_heads, self.head_dim) + v = v.view(total_seqlen, self.n_heads, self.head_dim) + + if self.rotary_emb is not None: + q, k = self.rotary_emb(q, k, cu_seqlens=cu_seqlens, max_seqlen=max_seq_len) + + kv = torch.stack([k, v], 1) + output = flash_attn_varlen_kvpacked_func( + q, kv, + cu_seqlens_q=cu_seqlens, cu_seqlens_k=cu_seqlens, + max_seqlen_q=max_seq_len, max_seqlen_k=max_seq_len, + dropout_p=0.0, causal=False, + ) + output = output.view(total_seqlen, self.n_heads * self.head_dim) + else: + # SDPA path (padded) + bsz, seqlen, _ = x.shape + q, k, v = torch.split(qkv, self.n_heads * self.head_dim, dim=-1) + q = q.view(bsz, seqlen, self.n_heads, self.head_dim) + k = k.view(bsz, seqlen, self.n_heads, self.head_dim) + v = v.view(bsz, seqlen, self.n_heads, self.head_dim) + + if self.rotary_emb is not None: + q, k = self.rotary_emb(q, k, position_ids=position_ids) + + q, k, v = q.transpose(1, 2), k.transpose(1, 2), v.transpose(1, 2) + + attn_mask = None + if attention_mask is not None: + attn_mask = attention_mask.unsqueeze(1).unsqueeze(2).bool() + + output = F.scaled_dot_product_attention(q, k, v, attn_mask=attn_mask, dropout_p=0.0, is_causal=False) + output = output.transpose(1, 2).contiguous().view(bsz, seqlen, self.n_heads * self.head_dim) + + return self.wo(output) + + +class FeedForward(nn.Module): + """SwiGLU feedforward network.""" + + def __init__(self, dim: int, hidden_mult: float = 4.0, multiple_of: int = 256, ffn_dim_multiplier: float = None): + super().__init__() + hidden_dim = int(2 * dim * hidden_mult / 3) + if ffn_dim_multiplier is not None: + hidden_dim = int(ffn_dim_multiplier * hidden_dim) + hidden_dim = multiple_of * ((hidden_dim + multiple_of - 1) // multiple_of) + self.w1 = nn.Linear(dim, hidden_dim, bias=False) + self.w2 = nn.Linear(hidden_dim, dim, bias=False) + self.w3 = nn.Linear(dim, hidden_dim, bias=False) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + return self.w2(swiglu(self.w1(x), self.w3(x))) + + +class TransformerBlock(nn.Module): + """Pre-norm transformer block.""" + + def __init__(self, dim: int, num_heads: int, norm_eps: float = 1e-6, + multiple_of: int = 256, ffn_dim_multiplier: float = None, use_rope: bool = True): + super().__init__() + self.attention = Attention(dim, num_heads, use_rope) + self.feed_forward = FeedForward(dim, multiple_of=multiple_of, ffn_dim_multiplier=ffn_dim_multiplier) + self.attention_norm = RMSNorm(dim, eps=norm_eps) + self.ffn_norm = RMSNorm(dim, eps=norm_eps) + + def forward( + self, + x: torch.Tensor, + cu_seqlens: Optional[torch.Tensor] = None, + max_seq_len: Optional[int] = None, + attention_mask: Optional[torch.Tensor] = None, + position_ids: Optional[torch.Tensor] = None, + ) -> torch.Tensor: + h = x + self.attention(self.attention_norm(x), cu_seqlens, max_seq_len, attention_mask, position_ids) + return h + self.feed_forward(self.ffn_norm(h)) + + +class TransformerLayers(nn.Module): + """Stack of transformer blocks with optional flash attention.""" + + def __init__(self, dim: int, num_heads: int, depth: int, norm_eps: float = 1e-6, + multiple_of: int = 256, ffn_dim_multiplier: float = None, use_rope: bool = True): + super().__init__() + self.dim = dim + self.layers = nn.ModuleList([ + TransformerBlock(dim, num_heads, norm_eps, multiple_of, ffn_dim_multiplier, use_rope) + for _ in range(depth) + ]) + self.gradient_checkpointing = False + + def forward(self, x: torch.Tensor, attention_mask: Optional[torch.Tensor] = None) -> torch.Tensor: + batch_size, seq_len = x.shape[:2] + cu_seqlens, max_seq_len_in_batch, indices, position_ids = None, None, None, None + + if FLASH_ATTN_AVAILABLE and attention_mask is not None and not attention_mask.all(): + x, indices, cu_seqlens, max_seq_len_in_batch, _ = unpad_input(x, attention_mask) + mask_for_layers = None + elif attention_mask is not None: + mask_long = attention_mask.long() + position_ids = (mask_long.cumsum(dim=1) - 1).clamp(min=0) + mask_for_layers = attention_mask + else: + mask_for_layers = None + + for layer in self.layers: + if self.training and self.gradient_checkpointing: + x = checkpoint(layer, x, cu_seqlens, max_seq_len_in_batch, mask_for_layers, position_ids, use_reentrant=False) + else: + x = layer(x, cu_seqlens, max_seq_len_in_batch, mask_for_layers, position_ids) + + if FLASH_ATTN_AVAILABLE and indices is not None: + x = pad_input(x, indices, batch_size, seq_len) + + return x + + +class GLM2Backbone(nn.Module): + """gLM2 protein language model backbone.""" + + def __init__(self, config: FlashPPIConfig): + super().__init__() + self.config = config + self.tok_embeddings = nn.Embedding(config.plm_vocab_size, config.plm_dim) + self.encoder = TransformerLayers( + dim=config.plm_dim, + num_heads=config.plm_heads, + depth=config.plm_depth, + norm_eps=config.plm_norm_eps, + multiple_of=config.plm_swiglu_multiple_of, + ffn_dim_multiplier=config.plm_ffn_dim_multiplier, + use_rope=True, + ) + + def forward(self, input_ids: torch.Tensor, attention_mask: Optional[torch.Tensor] = None) -> torch.Tensor: + h = self.tok_embeddings(input_ids) + return self.encoder(h, attention_mask) + + +class MLPHead(nn.Module): + """SwiGLU MLP projection head.""" + + def __init__(self, in_dim: int, out_dim: int, hidden_mult: float = 2.0): + super().__init__() + hidden_dim = int(in_dim * hidden_mult) + self.w1 = nn.Linear(in_dim, hidden_dim, bias=False) + self.w2 = nn.Linear(hidden_dim, out_dim, bias=False) + self.w3 = nn.Linear(in_dim, hidden_dim, bias=False) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + return self.w2(F.silu(self.w1(x)) * self.w3(x)) + + +class ContrastiveHead(nn.Module): + """CLIP-style contrastive head with mean pooling.""" + + def __init__(self, hidden_dim: int, embed_dim: int): + super().__init__() + self.head = MLPHead(hidden_dim, embed_dim) + + def forward(self, residue_embeds: torch.Tensor, mask: torch.Tensor) -> torch.Tensor: + mask = mask.unsqueeze(-1).bool() + embeds = torch.where(mask, residue_embeds, 0.0) + embeds = embeds.sum(dim=1) / mask.sum(dim=1).float().clamp(min=1.0) + return F.normalize(self.head(embeds), dim=-1) + + +class ContactHead(nn.Module): + """Contact prediction head using cross-attention between protein pairs.""" + + def __init__(self, input_dim: int, contact_dim: int, num_heads: int = 8, depth: int = 2): + super().__init__() + self.num_heads = num_heads + self.head_dim = contact_dim // num_heads + assert contact_dim % num_heads == 0 + + self.segment_embed = nn.Embedding(2, input_dim) + nn.init.normal_(self.segment_embed.weight, std=0.02) + + self.transformer = TransformerLayers(input_dim, num_heads, depth, use_rope=True) + self.norm = nn.LayerNorm(input_dim) + self.q_proj = nn.Linear(input_dim, contact_dim, bias=True) + self.k_proj = nn.Linear(input_dim, contact_dim, bias=True) + self.output_mix = nn.Linear(num_heads, 1) + nn.init.constant_(self.output_mix.bias, -3.0) + self.scale = self.head_dim ** -0.5 + + def forward( + self, + embed1: torch.Tensor, + embed2: torch.Tensor, + mask1: torch.Tensor, + mask2: torch.Tensor, + ) -> Tuple[torch.Tensor, torch.Tensor]: + B, L1, D = embed1.shape + _, L2, _ = embed2.shape + + seg1 = self.segment_embed(torch.zeros(L1, device=embed1.device, dtype=torch.long)) + seg2 = self.segment_embed(torch.ones(L2, device=embed1.device, dtype=torch.long)) + + x = torch.cat([embed1 + seg1.unsqueeze(0), embed2 + seg2.unsqueeze(0)], dim=1) + combined_mask = torch.cat([mask1, mask2], dim=1).bool() if mask1 is not None and mask2 is not None else None + + x = self.transformer(x, attention_mask=combined_mask) + + embed1 = self.norm(x[:, :L1, :]) + embed2 = self.norm(x[:, L1:, :]) + + q = self.q_proj(embed1).view(B, L1, self.num_heads, self.head_dim).transpose(1, 2) + k = self.k_proj(embed2).view(B, L2, self.num_heads, self.head_dim).transpose(1, 2) + + attn_logits = torch.matmul(q, k.transpose(-2, -1)) * self.scale + attn_logits = attn_logits.permute(0, 2, 3, 1).contiguous() + contact_logits = self.output_mix(attn_logits).squeeze(-1) + + if mask1 is not None and mask2 is not None: + valid_mask = (mask1.unsqueeze(2) * mask2.unsqueeze(1)).bool() + else: + valid_mask = torch.ones_like(contact_logits, dtype=torch.bool) + + return contact_logits, valid_mask + + +class FlashPPIPreTrainedModel(PreTrainedModel): + """Base class for FlashPPI models.""" + + config_class = FlashPPIConfig + base_model_prefix = "flashppi" + supports_gradient_checkpointing = True + + def _init_weights(self, module): + if isinstance(module, nn.Linear): + nn.init.normal_(module.weight, std=0.02) + if module.bias is not None: + nn.init.zeros_(module.bias) + elif isinstance(module, nn.Embedding): + nn.init.normal_(module.weight, std=0.02) + elif isinstance(module, RotaryEmbedding): + # Re-calculate the frequencies using the module's stored attributes + inv_freq = 1.0 / ( + module.base + ** ( + torch.arange(0, module.dim, 2, device=module.inv_freq.device, dtype=torch.float32) + / module.dim + ) + ) + # Force the buffer to update + with torch.no_grad(): + module.inv_freq.copy_(inv_freq) + +class FlashPPIModel(FlashPPIPreTrainedModel): + """FlashPPI model.""" + + def __init__(self, config: FlashPPIConfig): + super().__init__(config) + self.config = config + + # gLM2 backbone + self.plm = GLM2Backbone(config) + + # CLIP heads (asymmetric for query/key) + self.head_q = ContrastiveHead(config.plm_dim, config.clip_embed_dim) + self.head_k = ContrastiveHead(config.plm_dim, config.clip_embed_dim) + self.logit_scale = nn.Parameter(torch.ones([]) * 2.6593) # ln(1/0.07) + + # Contact prediction head + self.contact_head = ContactHead( + config.plm_dim, + config.contact_embed_dim, + num_heads=config.contact_num_heads, + depth=config.contact_transformer_depth, + ) + + self.post_init() + + def encode_protein( + self, + input_ids: torch.Tensor, + attention_mask: Optional[torch.Tensor] = None, + ) -> torch.Tensor: + """Encode a protein sequence to residue-level embeddings. + + Args: + input_ids: (B, L) token IDs from gLM2 tokenizer. + attention_mask: (B, L) attention mask. + + Returns: + (B, L, plm_dim) residue embeddings. + """ + return self.plm(input_ids, attention_mask) + + def predict_contacts( + self, + embed1: torch.Tensor, + embed2: torch.Tensor, + mask1: torch.Tensor, + mask2: torch.Tensor, + ) -> Tuple[torch.Tensor, torch.Tensor]: + """Predict contact map from pre-computed residue embeddings. + + This method is useful for efficient 2-stage inference where embeddings + are pre-computed and cached. + + Args: + embed1: (B, L1, D) residue embeddings for protein 1. + embed2: (B, L2, D) residue embeddings for protein 2. + mask1: (B, L1) attention mask for protein 1. + mask2: (B, L2) attention mask for protein 2. + + Returns: + contact_logits: (B, L1, L2) raw logits. + valid_mask: (B, L1, L2) mask for valid positions. + """ + return self.contact_head(embed1, embed2, mask1, mask2) + + def forward( + self, + input_ids1: torch.Tensor, + input_ids2: torch.Tensor, + attention_mask1: Optional[torch.Tensor] = None, + attention_mask2: Optional[torch.Tensor] = None, + return_dict: bool = True, + ) -> Union[Tuple, FlashPPIOutput]: + """Forward pass for protein pair interaction prediction. + + Args: + input_ids1: (B, L1) token IDs for protein 1. + input_ids2: (B, L2) token IDs for protein 2. + attention_mask1: (B, L1) attention mask for protein 1. + attention_mask2: (B, L2) attention mask for protein 2. + return_dict: Whether to return a FlashPPIOutput or tuple. + + Returns: + FlashPPIOutput with contact predictions and CLIP embeddings. + """ + B = input_ids1.shape[0] + L1, L2 = input_ids1.shape[1], input_ids2.shape[1] + + if attention_mask1 is None: + attention_mask1 = torch.ones_like(input_ids1) + if attention_mask2 is None: + attention_mask2 = torch.ones_like(input_ids2) + + # Encode both proteins in a single batched PLM call for efficiency + # Pad to same length if needed + if L1 != L2: + max_len = max(L1, L2) + if L1 < max_len: + pad_len = max_len - L1 + input_ids1 = F.pad(input_ids1, (0, pad_len), value=0) + attention_mask1 = F.pad(attention_mask1, (0, pad_len), value=0) + if L2 < max_len: + pad_len = max_len - L2 + input_ids2 = F.pad(input_ids2, (0, pad_len), value=0) + attention_mask2 = F.pad(attention_mask2, (0, pad_len), value=0) + + # Batch both sequences for single PLM forward pass + batched_input_ids = torch.cat([input_ids1, input_ids2], dim=0) + batched_attention_mask = torch.cat([attention_mask1, attention_mask2], dim=0) + batched_embeds = self.encode_protein(batched_input_ids, batched_attention_mask) + + # Split and trim back to original lengths + residue_embeds1 = batched_embeds[:B, :L1, :] + residue_embeds2 = batched_embeds[B:, :L2, :] + attention_mask1 = attention_mask1[:, :L1] + attention_mask2 = attention_mask2[:, :L2] + + # Contrastive embeddings + clip_embed1 = self.head_q(residue_embeds1, attention_mask1) + clip_embed2 = self.head_k(residue_embeds2, attention_mask2) + clip_score = (clip_embed1 * clip_embed2).sum(dim=-1) + + # Contact prediction + contact_logits, valid_mask = self.contact_head( + residue_embeds1, residue_embeds2, attention_mask1, attention_mask2 + ) + contact_map = torch.sigmoid(contact_logits) + + # Mask invalid positions before taking max + contact_map_masked = contact_map.masked_fill(~valid_mask, 0.0) + contact_score = contact_map_masked.flatten(1).max(dim=-1).values + + if not return_dict: + return (contact_map, contact_score, clip_embed1, clip_embed2, clip_score) + + return FlashPPIOutput( + contact_map=contact_map, + contact_score=contact_score, + clip_embed1=clip_embed1, + clip_embed2=clip_embed2, + clip_score=clip_score, + ) diff --git a/special_tokens_map.json b/special_tokens_map.json new file mode 100644 index 0000000000000000000000000000000000000000..cfbebd68549ad0845b37b122c8d80ebaa19f1f6a --- /dev/null +++ b/special_tokens_map.json @@ -0,0 +1,44 @@ +{ + "cls_token": { + "content": "", + "lstrip": false, + "normalized": false, + "rstrip": false, + "single_word": false + }, + "eos_token": { + "content": "", + "lstrip": false, + "normalized": false, + "rstrip": false, + "single_word": false + }, + "mask_token": { + "content": "", + "lstrip": false, + "normalized": false, + "rstrip": false, + "single_word": false + }, + "pad_token": { + "content": "", + "lstrip": false, + "normalized": false, + "rstrip": false, + "single_word": false + }, + "sep_token": { + "content": "", + "lstrip": false, + "normalized": false, + "rstrip": false, + "single_word": false + }, + "unk_token": { + "content": "", + "lstrip": false, + "normalized": false, + "rstrip": false, + "single_word": false + } +} diff --git a/tmp1non0gt5/__pycache__/_remote_module_non_scriptable.cpython-312.pyc b/tmp1non0gt5/__pycache__/_remote_module_non_scriptable.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..85c4aa6e55f1d13a8568f468e0bfdfa4c5354e52 Binary files /dev/null and b/tmp1non0gt5/__pycache__/_remote_module_non_scriptable.cpython-312.pyc differ diff --git a/tmp1non0gt5/_remote_module_non_scriptable.py b/tmp1non0gt5/_remote_module_non_scriptable.py new file mode 100644 index 0000000000000000000000000000000000000000..796e049370abda22e60acbeda5d65a4be33e1a87 --- /dev/null +++ b/tmp1non0gt5/_remote_module_non_scriptable.py @@ -0,0 +1,81 @@ +from typing import * + +import torch +import torch.distributed.rpc as rpc +from torch import Tensor +from torch._jit_internal import Future +from torch.distributed.rpc import RRef +from typing import Tuple # pyre-ignore: unused import + + +module_interface_cls = None + + +def forward_async(self, *args, **kwargs): + args = (self.module_rref, self.device, self.is_device_map_set, *args) + kwargs = {**kwargs} + return rpc.rpc_async( + self.module_rref.owner(), + _remote_forward, + args, + kwargs, + ) + + +def forward(self, *args, **kwargs): + args = (self.module_rref, self.device, self.is_device_map_set, *args) + kwargs = {**kwargs} + ret_fut = rpc.rpc_async( + self.module_rref.owner(), + _remote_forward, + args, + kwargs, + ) + return ret_fut.wait() + + +_generated_methods = [ + forward_async, + forward, +] + + + + +def _remote_forward( + module_rref: RRef[module_interface_cls], device: str, is_device_map_set: bool, *args, **kwargs): + module = module_rref.local_value() + device = torch.device(device) + + if device.type != "cuda": + return module.forward(*args, **kwargs) + + # If the module is on a cuda device, + # move any CPU tensor in args or kwargs to the same cuda device. + # Since torch script does not support generator expression, + # have to use concatenation instead of + # ``tuple(i.to(device) if isinstance(i, Tensor) else i for i in *args)``. + args = (*args,) + out_args: Tuple[()] = () + for arg in args: + arg = (arg.to(device),) if isinstance(arg, Tensor) else (arg,) + out_args = out_args + arg + + kwargs = {**kwargs} + for k, v in kwargs.items(): + if isinstance(v, Tensor): + kwargs[k] = kwargs[k].to(device) + + if is_device_map_set: + return module.forward(*out_args, **kwargs) + + # If the device map is empty, then only CPU tensors are allowed to send over wire, + # so have to move any GPU tensor to CPU in the output. + # Since torch script does not support generator expression, + # have to use concatenation instead of + # ``tuple(i.cpu() if isinstance(i, Tensor) else i for i in module.forward(*out_args, **kwargs))``. + ret: Tuple[()] = () + for i in module.forward(*out_args, **kwargs): + i = (i.cpu(),) if isinstance(i, Tensor) else (i,) + ret = ret + i + return ret diff --git a/tmp38o1kqky/__pycache__/_remote_module_non_scriptable.cpython-312.pyc b/tmp38o1kqky/__pycache__/_remote_module_non_scriptable.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..96b467bf0e41921193eed8ce318353a3dfb88b48 Binary files /dev/null and b/tmp38o1kqky/__pycache__/_remote_module_non_scriptable.cpython-312.pyc differ diff --git a/tmp38o1kqky/_remote_module_non_scriptable.py b/tmp38o1kqky/_remote_module_non_scriptable.py new file mode 100644 index 0000000000000000000000000000000000000000..796e049370abda22e60acbeda5d65a4be33e1a87 --- /dev/null +++ b/tmp38o1kqky/_remote_module_non_scriptable.py @@ -0,0 +1,81 @@ +from typing import * + +import torch +import torch.distributed.rpc as rpc +from torch import Tensor +from torch._jit_internal import Future +from torch.distributed.rpc import RRef +from typing import Tuple # pyre-ignore: unused import + + +module_interface_cls = None + + +def forward_async(self, *args, **kwargs): + args = (self.module_rref, self.device, self.is_device_map_set, *args) + kwargs = {**kwargs} + return rpc.rpc_async( + self.module_rref.owner(), + _remote_forward, + args, + kwargs, + ) + + +def forward(self, *args, **kwargs): + args = (self.module_rref, self.device, self.is_device_map_set, *args) + kwargs = {**kwargs} + ret_fut = rpc.rpc_async( + self.module_rref.owner(), + _remote_forward, + args, + kwargs, + ) + return ret_fut.wait() + + +_generated_methods = [ + forward_async, + forward, +] + + + + +def _remote_forward( + module_rref: RRef[module_interface_cls], device: str, is_device_map_set: bool, *args, **kwargs): + module = module_rref.local_value() + device = torch.device(device) + + if device.type != "cuda": + return module.forward(*args, **kwargs) + + # If the module is on a cuda device, + # move any CPU tensor in args or kwargs to the same cuda device. + # Since torch script does not support generator expression, + # have to use concatenation instead of + # ``tuple(i.to(device) if isinstance(i, Tensor) else i for i in *args)``. + args = (*args,) + out_args: Tuple[()] = () + for arg in args: + arg = (arg.to(device),) if isinstance(arg, Tensor) else (arg,) + out_args = out_args + arg + + kwargs = {**kwargs} + for k, v in kwargs.items(): + if isinstance(v, Tensor): + kwargs[k] = kwargs[k].to(device) + + if is_device_map_set: + return module.forward(*out_args, **kwargs) + + # If the device map is empty, then only CPU tensors are allowed to send over wire, + # so have to move any GPU tensor to CPU in the output. + # Since torch script does not support generator expression, + # have to use concatenation instead of + # ``tuple(i.cpu() if isinstance(i, Tensor) else i for i in module.forward(*out_args, **kwargs))``. + ret: Tuple[()] = () + for i in module.forward(*out_args, **kwargs): + i = (i.cpu(),) if isinstance(i, Tensor) else (i,) + ret = ret + i + return ret diff --git a/tmp69csjip6/__pycache__/_remote_module_non_scriptable.cpython-312.pyc b/tmp69csjip6/__pycache__/_remote_module_non_scriptable.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..867737ba0f81526af9787a8abf43a309b2d98fc3 Binary files /dev/null and b/tmp69csjip6/__pycache__/_remote_module_non_scriptable.cpython-312.pyc differ diff --git a/tmp69csjip6/_remote_module_non_scriptable.py b/tmp69csjip6/_remote_module_non_scriptable.py new file mode 100644 index 0000000000000000000000000000000000000000..796e049370abda22e60acbeda5d65a4be33e1a87 --- /dev/null +++ b/tmp69csjip6/_remote_module_non_scriptable.py @@ -0,0 +1,81 @@ +from typing import * + +import torch +import torch.distributed.rpc as rpc +from torch import Tensor +from torch._jit_internal import Future +from torch.distributed.rpc import RRef +from typing import Tuple # pyre-ignore: unused import + + +module_interface_cls = None + + +def forward_async(self, *args, **kwargs): + args = (self.module_rref, self.device, self.is_device_map_set, *args) + kwargs = {**kwargs} + return rpc.rpc_async( + self.module_rref.owner(), + _remote_forward, + args, + kwargs, + ) + + +def forward(self, *args, **kwargs): + args = (self.module_rref, self.device, self.is_device_map_set, *args) + kwargs = {**kwargs} + ret_fut = rpc.rpc_async( + self.module_rref.owner(), + _remote_forward, + args, + kwargs, + ) + return ret_fut.wait() + + +_generated_methods = [ + forward_async, + forward, +] + + + + +def _remote_forward( + module_rref: RRef[module_interface_cls], device: str, is_device_map_set: bool, *args, **kwargs): + module = module_rref.local_value() + device = torch.device(device) + + if device.type != "cuda": + return module.forward(*args, **kwargs) + + # If the module is on a cuda device, + # move any CPU tensor in args or kwargs to the same cuda device. + # Since torch script does not support generator expression, + # have to use concatenation instead of + # ``tuple(i.to(device) if isinstance(i, Tensor) else i for i in *args)``. + args = (*args,) + out_args: Tuple[()] = () + for arg in args: + arg = (arg.to(device),) if isinstance(arg, Tensor) else (arg,) + out_args = out_args + arg + + kwargs = {**kwargs} + for k, v in kwargs.items(): + if isinstance(v, Tensor): + kwargs[k] = kwargs[k].to(device) + + if is_device_map_set: + return module.forward(*out_args, **kwargs) + + # If the device map is empty, then only CPU tensors are allowed to send over wire, + # so have to move any GPU tensor to CPU in the output. + # Since torch script does not support generator expression, + # have to use concatenation instead of + # ``tuple(i.cpu() if isinstance(i, Tensor) else i for i in module.forward(*out_args, **kwargs))``. + ret: Tuple[()] = () + for i in module.forward(*out_args, **kwargs): + i = (i.cpu(),) if isinstance(i, Tensor) else (i,) + ret = ret + i + return ret diff --git a/tmp9hzp29qv/__pycache__/_remote_module_non_scriptable.cpython-312.pyc b/tmp9hzp29qv/__pycache__/_remote_module_non_scriptable.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..1b812a0f789e8e27d88011f58b22bc556027f798 Binary files /dev/null and b/tmp9hzp29qv/__pycache__/_remote_module_non_scriptable.cpython-312.pyc differ diff --git a/tmp9hzp29qv/_remote_module_non_scriptable.py b/tmp9hzp29qv/_remote_module_non_scriptable.py new file mode 100644 index 0000000000000000000000000000000000000000..796e049370abda22e60acbeda5d65a4be33e1a87 --- /dev/null +++ b/tmp9hzp29qv/_remote_module_non_scriptable.py @@ -0,0 +1,81 @@ +from typing import * + +import torch +import torch.distributed.rpc as rpc +from torch import Tensor +from torch._jit_internal import Future +from torch.distributed.rpc import RRef +from typing import Tuple # pyre-ignore: unused import + + +module_interface_cls = None + + +def forward_async(self, *args, **kwargs): + args = (self.module_rref, self.device, self.is_device_map_set, *args) + kwargs = {**kwargs} + return rpc.rpc_async( + self.module_rref.owner(), + _remote_forward, + args, + kwargs, + ) + + +def forward(self, *args, **kwargs): + args = (self.module_rref, self.device, self.is_device_map_set, *args) + kwargs = {**kwargs} + ret_fut = rpc.rpc_async( + self.module_rref.owner(), + _remote_forward, + args, + kwargs, + ) + return ret_fut.wait() + + +_generated_methods = [ + forward_async, + forward, +] + + + + +def _remote_forward( + module_rref: RRef[module_interface_cls], device: str, is_device_map_set: bool, *args, **kwargs): + module = module_rref.local_value() + device = torch.device(device) + + if device.type != "cuda": + return module.forward(*args, **kwargs) + + # If the module is on a cuda device, + # move any CPU tensor in args or kwargs to the same cuda device. + # Since torch script does not support generator expression, + # have to use concatenation instead of + # ``tuple(i.to(device) if isinstance(i, Tensor) else i for i in *args)``. + args = (*args,) + out_args: Tuple[()] = () + for arg in args: + arg = (arg.to(device),) if isinstance(arg, Tensor) else (arg,) + out_args = out_args + arg + + kwargs = {**kwargs} + for k, v in kwargs.items(): + if isinstance(v, Tensor): + kwargs[k] = kwargs[k].to(device) + + if is_device_map_set: + return module.forward(*out_args, **kwargs) + + # If the device map is empty, then only CPU tensors are allowed to send over wire, + # so have to move any GPU tensor to CPU in the output. + # Since torch script does not support generator expression, + # have to use concatenation instead of + # ``tuple(i.cpu() if isinstance(i, Tensor) else i for i in module.forward(*out_args, **kwargs))``. + ret: Tuple[()] = () + for i in module.forward(*out_args, **kwargs): + i = (i.cpu(),) if isinstance(i, Tensor) else (i,) + ret = ret + i + return ret diff --git a/tmp_j0mzuvp/__pycache__/_remote_module_non_scriptable.cpython-312.pyc b/tmp_j0mzuvp/__pycache__/_remote_module_non_scriptable.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..070068eeed0019a53556d6f67ebc03ca47710b29 Binary files /dev/null and b/tmp_j0mzuvp/__pycache__/_remote_module_non_scriptable.cpython-312.pyc differ diff --git a/tmp_j0mzuvp/_remote_module_non_scriptable.py b/tmp_j0mzuvp/_remote_module_non_scriptable.py new file mode 100644 index 0000000000000000000000000000000000000000..796e049370abda22e60acbeda5d65a4be33e1a87 --- /dev/null +++ b/tmp_j0mzuvp/_remote_module_non_scriptable.py @@ -0,0 +1,81 @@ +from typing import * + +import torch +import torch.distributed.rpc as rpc +from torch import Tensor +from torch._jit_internal import Future +from torch.distributed.rpc import RRef +from typing import Tuple # pyre-ignore: unused import + + +module_interface_cls = None + + +def forward_async(self, *args, **kwargs): + args = (self.module_rref, self.device, self.is_device_map_set, *args) + kwargs = {**kwargs} + return rpc.rpc_async( + self.module_rref.owner(), + _remote_forward, + args, + kwargs, + ) + + +def forward(self, *args, **kwargs): + args = (self.module_rref, self.device, self.is_device_map_set, *args) + kwargs = {**kwargs} + ret_fut = rpc.rpc_async( + self.module_rref.owner(), + _remote_forward, + args, + kwargs, + ) + return ret_fut.wait() + + +_generated_methods = [ + forward_async, + forward, +] + + + + +def _remote_forward( + module_rref: RRef[module_interface_cls], device: str, is_device_map_set: bool, *args, **kwargs): + module = module_rref.local_value() + device = torch.device(device) + + if device.type != "cuda": + return module.forward(*args, **kwargs) + + # If the module is on a cuda device, + # move any CPU tensor in args or kwargs to the same cuda device. + # Since torch script does not support generator expression, + # have to use concatenation instead of + # ``tuple(i.to(device) if isinstance(i, Tensor) else i for i in *args)``. + args = (*args,) + out_args: Tuple[()] = () + for arg in args: + arg = (arg.to(device),) if isinstance(arg, Tensor) else (arg,) + out_args = out_args + arg + + kwargs = {**kwargs} + for k, v in kwargs.items(): + if isinstance(v, Tensor): + kwargs[k] = kwargs[k].to(device) + + if is_device_map_set: + return module.forward(*out_args, **kwargs) + + # If the device map is empty, then only CPU tensors are allowed to send over wire, + # so have to move any GPU tensor to CPU in the output. + # Since torch script does not support generator expression, + # have to use concatenation instead of + # ``tuple(i.cpu() if isinstance(i, Tensor) else i for i in module.forward(*out_args, **kwargs))``. + ret: Tuple[()] = () + for i in module.forward(*out_args, **kwargs): + i = (i.cpu(),) if isinstance(i, Tensor) else (i,) + ret = ret + i + return ret diff --git a/tmpd760ylpb/__pycache__/_remote_module_non_scriptable.cpython-312.pyc b/tmpd760ylpb/__pycache__/_remote_module_non_scriptable.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..e72263f9e6cf8cd55bc4bebbbac0cf326d9811e6 Binary files /dev/null and b/tmpd760ylpb/__pycache__/_remote_module_non_scriptable.cpython-312.pyc differ diff --git a/tmpd760ylpb/_remote_module_non_scriptable.py b/tmpd760ylpb/_remote_module_non_scriptable.py new file mode 100644 index 0000000000000000000000000000000000000000..796e049370abda22e60acbeda5d65a4be33e1a87 --- /dev/null +++ b/tmpd760ylpb/_remote_module_non_scriptable.py @@ -0,0 +1,81 @@ +from typing import * + +import torch +import torch.distributed.rpc as rpc +from torch import Tensor +from torch._jit_internal import Future +from torch.distributed.rpc import RRef +from typing import Tuple # pyre-ignore: unused import + + +module_interface_cls = None + + +def forward_async(self, *args, **kwargs): + args = (self.module_rref, self.device, self.is_device_map_set, *args) + kwargs = {**kwargs} + return rpc.rpc_async( + self.module_rref.owner(), + _remote_forward, + args, + kwargs, + ) + + +def forward(self, *args, **kwargs): + args = (self.module_rref, self.device, self.is_device_map_set, *args) + kwargs = {**kwargs} + ret_fut = rpc.rpc_async( + self.module_rref.owner(), + _remote_forward, + args, + kwargs, + ) + return ret_fut.wait() + + +_generated_methods = [ + forward_async, + forward, +] + + + + +def _remote_forward( + module_rref: RRef[module_interface_cls], device: str, is_device_map_set: bool, *args, **kwargs): + module = module_rref.local_value() + device = torch.device(device) + + if device.type != "cuda": + return module.forward(*args, **kwargs) + + # If the module is on a cuda device, + # move any CPU tensor in args or kwargs to the same cuda device. + # Since torch script does not support generator expression, + # have to use concatenation instead of + # ``tuple(i.to(device) if isinstance(i, Tensor) else i for i in *args)``. + args = (*args,) + out_args: Tuple[()] = () + for arg in args: + arg = (arg.to(device),) if isinstance(arg, Tensor) else (arg,) + out_args = out_args + arg + + kwargs = {**kwargs} + for k, v in kwargs.items(): + if isinstance(v, Tensor): + kwargs[k] = kwargs[k].to(device) + + if is_device_map_set: + return module.forward(*out_args, **kwargs) + + # If the device map is empty, then only CPU tensors are allowed to send over wire, + # so have to move any GPU tensor to CPU in the output. + # Since torch script does not support generator expression, + # have to use concatenation instead of + # ``tuple(i.cpu() if isinstance(i, Tensor) else i for i in module.forward(*out_args, **kwargs))``. + ret: Tuple[()] = () + for i in module.forward(*out_args, **kwargs): + i = (i.cpu(),) if isinstance(i, Tensor) else (i,) + ret = ret + i + return ret diff --git a/tmpdxfzs37c/__pycache__/_remote_module_non_scriptable.cpython-312.pyc b/tmpdxfzs37c/__pycache__/_remote_module_non_scriptable.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..ca45c2af2a0c171da2e15691452992b511d46bb5 Binary files /dev/null and b/tmpdxfzs37c/__pycache__/_remote_module_non_scriptable.cpython-312.pyc differ diff --git a/tmpdxfzs37c/_remote_module_non_scriptable.py b/tmpdxfzs37c/_remote_module_non_scriptable.py new file mode 100644 index 0000000000000000000000000000000000000000..796e049370abda22e60acbeda5d65a4be33e1a87 --- /dev/null +++ b/tmpdxfzs37c/_remote_module_non_scriptable.py @@ -0,0 +1,81 @@ +from typing import * + +import torch +import torch.distributed.rpc as rpc +from torch import Tensor +from torch._jit_internal import Future +from torch.distributed.rpc import RRef +from typing import Tuple # pyre-ignore: unused import + + +module_interface_cls = None + + +def forward_async(self, *args, **kwargs): + args = (self.module_rref, self.device, self.is_device_map_set, *args) + kwargs = {**kwargs} + return rpc.rpc_async( + self.module_rref.owner(), + _remote_forward, + args, + kwargs, + ) + + +def forward(self, *args, **kwargs): + args = (self.module_rref, self.device, self.is_device_map_set, *args) + kwargs = {**kwargs} + ret_fut = rpc.rpc_async( + self.module_rref.owner(), + _remote_forward, + args, + kwargs, + ) + return ret_fut.wait() + + +_generated_methods = [ + forward_async, + forward, +] + + + + +def _remote_forward( + module_rref: RRef[module_interface_cls], device: str, is_device_map_set: bool, *args, **kwargs): + module = module_rref.local_value() + device = torch.device(device) + + if device.type != "cuda": + return module.forward(*args, **kwargs) + + # If the module is on a cuda device, + # move any CPU tensor in args or kwargs to the same cuda device. + # Since torch script does not support generator expression, + # have to use concatenation instead of + # ``tuple(i.to(device) if isinstance(i, Tensor) else i for i in *args)``. + args = (*args,) + out_args: Tuple[()] = () + for arg in args: + arg = (arg.to(device),) if isinstance(arg, Tensor) else (arg,) + out_args = out_args + arg + + kwargs = {**kwargs} + for k, v in kwargs.items(): + if isinstance(v, Tensor): + kwargs[k] = kwargs[k].to(device) + + if is_device_map_set: + return module.forward(*out_args, **kwargs) + + # If the device map is empty, then only CPU tensors are allowed to send over wire, + # so have to move any GPU tensor to CPU in the output. + # Since torch script does not support generator expression, + # have to use concatenation instead of + # ``tuple(i.cpu() if isinstance(i, Tensor) else i for i in module.forward(*out_args, **kwargs))``. + ret: Tuple[()] = () + for i in module.forward(*out_args, **kwargs): + i = (i.cpu(),) if isinstance(i, Tensor) else (i,) + ret = ret + i + return ret diff --git a/tmpf5f7c_r3/__pycache__/_remote_module_non_scriptable.cpython-312.pyc b/tmpf5f7c_r3/__pycache__/_remote_module_non_scriptable.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..325dc973fdbfe6526ebf71f7134384de8d9fc914 Binary files /dev/null and b/tmpf5f7c_r3/__pycache__/_remote_module_non_scriptable.cpython-312.pyc differ diff --git a/tmpf5f7c_r3/_remote_module_non_scriptable.py b/tmpf5f7c_r3/_remote_module_non_scriptable.py new file mode 100644 index 0000000000000000000000000000000000000000..796e049370abda22e60acbeda5d65a4be33e1a87 --- /dev/null +++ b/tmpf5f7c_r3/_remote_module_non_scriptable.py @@ -0,0 +1,81 @@ +from typing import * + +import torch +import torch.distributed.rpc as rpc +from torch import Tensor +from torch._jit_internal import Future +from torch.distributed.rpc import RRef +from typing import Tuple # pyre-ignore: unused import + + +module_interface_cls = None + + +def forward_async(self, *args, **kwargs): + args = (self.module_rref, self.device, self.is_device_map_set, *args) + kwargs = {**kwargs} + return rpc.rpc_async( + self.module_rref.owner(), + _remote_forward, + args, + kwargs, + ) + + +def forward(self, *args, **kwargs): + args = (self.module_rref, self.device, self.is_device_map_set, *args) + kwargs = {**kwargs} + ret_fut = rpc.rpc_async( + self.module_rref.owner(), + _remote_forward, + args, + kwargs, + ) + return ret_fut.wait() + + +_generated_methods = [ + forward_async, + forward, +] + + + + +def _remote_forward( + module_rref: RRef[module_interface_cls], device: str, is_device_map_set: bool, *args, **kwargs): + module = module_rref.local_value() + device = torch.device(device) + + if device.type != "cuda": + return module.forward(*args, **kwargs) + + # If the module is on a cuda device, + # move any CPU tensor in args or kwargs to the same cuda device. + # Since torch script does not support generator expression, + # have to use concatenation instead of + # ``tuple(i.to(device) if isinstance(i, Tensor) else i for i in *args)``. + args = (*args,) + out_args: Tuple[()] = () + for arg in args: + arg = (arg.to(device),) if isinstance(arg, Tensor) else (arg,) + out_args = out_args + arg + + kwargs = {**kwargs} + for k, v in kwargs.items(): + if isinstance(v, Tensor): + kwargs[k] = kwargs[k].to(device) + + if is_device_map_set: + return module.forward(*out_args, **kwargs) + + # If the device map is empty, then only CPU tensors are allowed to send over wire, + # so have to move any GPU tensor to CPU in the output. + # Since torch script does not support generator expression, + # have to use concatenation instead of + # ``tuple(i.cpu() if isinstance(i, Tensor) else i for i in module.forward(*out_args, **kwargs))``. + ret: Tuple[()] = () + for i in module.forward(*out_args, **kwargs): + i = (i.cpu(),) if isinstance(i, Tensor) else (i,) + ret = ret + i + return ret diff --git a/tmpiw5dwlcx/__pycache__/_remote_module_non_scriptable.cpython-312.pyc b/tmpiw5dwlcx/__pycache__/_remote_module_non_scriptable.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..646b8655780662bedd157922b70a2aac21598f85 Binary files /dev/null and b/tmpiw5dwlcx/__pycache__/_remote_module_non_scriptable.cpython-312.pyc differ diff --git a/tmpiw5dwlcx/_remote_module_non_scriptable.py b/tmpiw5dwlcx/_remote_module_non_scriptable.py new file mode 100644 index 0000000000000000000000000000000000000000..796e049370abda22e60acbeda5d65a4be33e1a87 --- /dev/null +++ b/tmpiw5dwlcx/_remote_module_non_scriptable.py @@ -0,0 +1,81 @@ +from typing import * + +import torch +import torch.distributed.rpc as rpc +from torch import Tensor +from torch._jit_internal import Future +from torch.distributed.rpc import RRef +from typing import Tuple # pyre-ignore: unused import + + +module_interface_cls = None + + +def forward_async(self, *args, **kwargs): + args = (self.module_rref, self.device, self.is_device_map_set, *args) + kwargs = {**kwargs} + return rpc.rpc_async( + self.module_rref.owner(), + _remote_forward, + args, + kwargs, + ) + + +def forward(self, *args, **kwargs): + args = (self.module_rref, self.device, self.is_device_map_set, *args) + kwargs = {**kwargs} + ret_fut = rpc.rpc_async( + self.module_rref.owner(), + _remote_forward, + args, + kwargs, + ) + return ret_fut.wait() + + +_generated_methods = [ + forward_async, + forward, +] + + + + +def _remote_forward( + module_rref: RRef[module_interface_cls], device: str, is_device_map_set: bool, *args, **kwargs): + module = module_rref.local_value() + device = torch.device(device) + + if device.type != "cuda": + return module.forward(*args, **kwargs) + + # If the module is on a cuda device, + # move any CPU tensor in args or kwargs to the same cuda device. + # Since torch script does not support generator expression, + # have to use concatenation instead of + # ``tuple(i.to(device) if isinstance(i, Tensor) else i for i in *args)``. + args = (*args,) + out_args: Tuple[()] = () + for arg in args: + arg = (arg.to(device),) if isinstance(arg, Tensor) else (arg,) + out_args = out_args + arg + + kwargs = {**kwargs} + for k, v in kwargs.items(): + if isinstance(v, Tensor): + kwargs[k] = kwargs[k].to(device) + + if is_device_map_set: + return module.forward(*out_args, **kwargs) + + # If the device map is empty, then only CPU tensors are allowed to send over wire, + # so have to move any GPU tensor to CPU in the output. + # Since torch script does not support generator expression, + # have to use concatenation instead of + # ``tuple(i.cpu() if isinstance(i, Tensor) else i for i in module.forward(*out_args, **kwargs))``. + ret: Tuple[()] = () + for i in module.forward(*out_args, **kwargs): + i = (i.cpu(),) if isinstance(i, Tensor) else (i,) + ret = ret + i + return ret diff --git a/tmpnkix_nz4/__pycache__/_remote_module_non_scriptable.cpython-312.pyc b/tmpnkix_nz4/__pycache__/_remote_module_non_scriptable.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..082e7ec4a6d16109343b5be79c2a78698532d74f Binary files /dev/null and b/tmpnkix_nz4/__pycache__/_remote_module_non_scriptable.cpython-312.pyc differ diff --git a/tmpnkix_nz4/_remote_module_non_scriptable.py b/tmpnkix_nz4/_remote_module_non_scriptable.py new file mode 100644 index 0000000000000000000000000000000000000000..796e049370abda22e60acbeda5d65a4be33e1a87 --- /dev/null +++ b/tmpnkix_nz4/_remote_module_non_scriptable.py @@ -0,0 +1,81 @@ +from typing import * + +import torch +import torch.distributed.rpc as rpc +from torch import Tensor +from torch._jit_internal import Future +from torch.distributed.rpc import RRef +from typing import Tuple # pyre-ignore: unused import + + +module_interface_cls = None + + +def forward_async(self, *args, **kwargs): + args = (self.module_rref, self.device, self.is_device_map_set, *args) + kwargs = {**kwargs} + return rpc.rpc_async( + self.module_rref.owner(), + _remote_forward, + args, + kwargs, + ) + + +def forward(self, *args, **kwargs): + args = (self.module_rref, self.device, self.is_device_map_set, *args) + kwargs = {**kwargs} + ret_fut = rpc.rpc_async( + self.module_rref.owner(), + _remote_forward, + args, + kwargs, + ) + return ret_fut.wait() + + +_generated_methods = [ + forward_async, + forward, +] + + + + +def _remote_forward( + module_rref: RRef[module_interface_cls], device: str, is_device_map_set: bool, *args, **kwargs): + module = module_rref.local_value() + device = torch.device(device) + + if device.type != "cuda": + return module.forward(*args, **kwargs) + + # If the module is on a cuda device, + # move any CPU tensor in args or kwargs to the same cuda device. + # Since torch script does not support generator expression, + # have to use concatenation instead of + # ``tuple(i.to(device) if isinstance(i, Tensor) else i for i in *args)``. + args = (*args,) + out_args: Tuple[()] = () + for arg in args: + arg = (arg.to(device),) if isinstance(arg, Tensor) else (arg,) + out_args = out_args + arg + + kwargs = {**kwargs} + for k, v in kwargs.items(): + if isinstance(v, Tensor): + kwargs[k] = kwargs[k].to(device) + + if is_device_map_set: + return module.forward(*out_args, **kwargs) + + # If the device map is empty, then only CPU tensors are allowed to send over wire, + # so have to move any GPU tensor to CPU in the output. + # Since torch script does not support generator expression, + # have to use concatenation instead of + # ``tuple(i.cpu() if isinstance(i, Tensor) else i for i in module.forward(*out_args, **kwargs))``. + ret: Tuple[()] = () + for i in module.forward(*out_args, **kwargs): + i = (i.cpu(),) if isinstance(i, Tensor) else (i,) + ret = ret + i + return ret diff --git a/tmpp33g8q8r/__pycache__/_remote_module_non_scriptable.cpython-312.pyc b/tmpp33g8q8r/__pycache__/_remote_module_non_scriptable.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..b57ba7cba5a15322ddcac7045f9e426438eb5792 Binary files /dev/null and b/tmpp33g8q8r/__pycache__/_remote_module_non_scriptable.cpython-312.pyc differ diff --git a/tmpp33g8q8r/_remote_module_non_scriptable.py b/tmpp33g8q8r/_remote_module_non_scriptable.py new file mode 100644 index 0000000000000000000000000000000000000000..796e049370abda22e60acbeda5d65a4be33e1a87 --- /dev/null +++ b/tmpp33g8q8r/_remote_module_non_scriptable.py @@ -0,0 +1,81 @@ +from typing import * + +import torch +import torch.distributed.rpc as rpc +from torch import Tensor +from torch._jit_internal import Future +from torch.distributed.rpc import RRef +from typing import Tuple # pyre-ignore: unused import + + +module_interface_cls = None + + +def forward_async(self, *args, **kwargs): + args = (self.module_rref, self.device, self.is_device_map_set, *args) + kwargs = {**kwargs} + return rpc.rpc_async( + self.module_rref.owner(), + _remote_forward, + args, + kwargs, + ) + + +def forward(self, *args, **kwargs): + args = (self.module_rref, self.device, self.is_device_map_set, *args) + kwargs = {**kwargs} + ret_fut = rpc.rpc_async( + self.module_rref.owner(), + _remote_forward, + args, + kwargs, + ) + return ret_fut.wait() + + +_generated_methods = [ + forward_async, + forward, +] + + + + +def _remote_forward( + module_rref: RRef[module_interface_cls], device: str, is_device_map_set: bool, *args, **kwargs): + module = module_rref.local_value() + device = torch.device(device) + + if device.type != "cuda": + return module.forward(*args, **kwargs) + + # If the module is on a cuda device, + # move any CPU tensor in args or kwargs to the same cuda device. + # Since torch script does not support generator expression, + # have to use concatenation instead of + # ``tuple(i.to(device) if isinstance(i, Tensor) else i for i in *args)``. + args = (*args,) + out_args: Tuple[()] = () + for arg in args: + arg = (arg.to(device),) if isinstance(arg, Tensor) else (arg,) + out_args = out_args + arg + + kwargs = {**kwargs} + for k, v in kwargs.items(): + if isinstance(v, Tensor): + kwargs[k] = kwargs[k].to(device) + + if is_device_map_set: + return module.forward(*out_args, **kwargs) + + # If the device map is empty, then only CPU tensors are allowed to send over wire, + # so have to move any GPU tensor to CPU in the output. + # Since torch script does not support generator expression, + # have to use concatenation instead of + # ``tuple(i.cpu() if isinstance(i, Tensor) else i for i in module.forward(*out_args, **kwargs))``. + ret: Tuple[()] = () + for i in module.forward(*out_args, **kwargs): + i = (i.cpu(),) if isinstance(i, Tensor) else (i,) + ret = ret + i + return ret diff --git a/tmppbcdn8h9/__pycache__/_remote_module_non_scriptable.cpython-312.pyc b/tmppbcdn8h9/__pycache__/_remote_module_non_scriptable.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..d298ee8fc822ccfa948ac37346c977d0d438f064 Binary files /dev/null and b/tmppbcdn8h9/__pycache__/_remote_module_non_scriptable.cpython-312.pyc differ diff --git a/tmppbcdn8h9/_remote_module_non_scriptable.py b/tmppbcdn8h9/_remote_module_non_scriptable.py new file mode 100644 index 0000000000000000000000000000000000000000..796e049370abda22e60acbeda5d65a4be33e1a87 --- /dev/null +++ b/tmppbcdn8h9/_remote_module_non_scriptable.py @@ -0,0 +1,81 @@ +from typing import * + +import torch +import torch.distributed.rpc as rpc +from torch import Tensor +from torch._jit_internal import Future +from torch.distributed.rpc import RRef +from typing import Tuple # pyre-ignore: unused import + + +module_interface_cls = None + + +def forward_async(self, *args, **kwargs): + args = (self.module_rref, self.device, self.is_device_map_set, *args) + kwargs = {**kwargs} + return rpc.rpc_async( + self.module_rref.owner(), + _remote_forward, + args, + kwargs, + ) + + +def forward(self, *args, **kwargs): + args = (self.module_rref, self.device, self.is_device_map_set, *args) + kwargs = {**kwargs} + ret_fut = rpc.rpc_async( + self.module_rref.owner(), + _remote_forward, + args, + kwargs, + ) + return ret_fut.wait() + + +_generated_methods = [ + forward_async, + forward, +] + + + + +def _remote_forward( + module_rref: RRef[module_interface_cls], device: str, is_device_map_set: bool, *args, **kwargs): + module = module_rref.local_value() + device = torch.device(device) + + if device.type != "cuda": + return module.forward(*args, **kwargs) + + # If the module is on a cuda device, + # move any CPU tensor in args or kwargs to the same cuda device. + # Since torch script does not support generator expression, + # have to use concatenation instead of + # ``tuple(i.to(device) if isinstance(i, Tensor) else i for i in *args)``. + args = (*args,) + out_args: Tuple[()] = () + for arg in args: + arg = (arg.to(device),) if isinstance(arg, Tensor) else (arg,) + out_args = out_args + arg + + kwargs = {**kwargs} + for k, v in kwargs.items(): + if isinstance(v, Tensor): + kwargs[k] = kwargs[k].to(device) + + if is_device_map_set: + return module.forward(*out_args, **kwargs) + + # If the device map is empty, then only CPU tensors are allowed to send over wire, + # so have to move any GPU tensor to CPU in the output. + # Since torch script does not support generator expression, + # have to use concatenation instead of + # ``tuple(i.cpu() if isinstance(i, Tensor) else i for i in module.forward(*out_args, **kwargs))``. + ret: Tuple[()] = () + for i in module.forward(*out_args, **kwargs): + i = (i.cpu(),) if isinstance(i, Tensor) else (i,) + ret = ret + i + return ret diff --git a/tmpqiqrdevn/__pycache__/_remote_module_non_scriptable.cpython-312.pyc b/tmpqiqrdevn/__pycache__/_remote_module_non_scriptable.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..df02dd0ceeb4701e6a7a50f1a075caeac18113f2 Binary files /dev/null and b/tmpqiqrdevn/__pycache__/_remote_module_non_scriptable.cpython-312.pyc differ diff --git a/tmpqiqrdevn/_remote_module_non_scriptable.py b/tmpqiqrdevn/_remote_module_non_scriptable.py new file mode 100644 index 0000000000000000000000000000000000000000..796e049370abda22e60acbeda5d65a4be33e1a87 --- /dev/null +++ b/tmpqiqrdevn/_remote_module_non_scriptable.py @@ -0,0 +1,81 @@ +from typing import * + +import torch +import torch.distributed.rpc as rpc +from torch import Tensor +from torch._jit_internal import Future +from torch.distributed.rpc import RRef +from typing import Tuple # pyre-ignore: unused import + + +module_interface_cls = None + + +def forward_async(self, *args, **kwargs): + args = (self.module_rref, self.device, self.is_device_map_set, *args) + kwargs = {**kwargs} + return rpc.rpc_async( + self.module_rref.owner(), + _remote_forward, + args, + kwargs, + ) + + +def forward(self, *args, **kwargs): + args = (self.module_rref, self.device, self.is_device_map_set, *args) + kwargs = {**kwargs} + ret_fut = rpc.rpc_async( + self.module_rref.owner(), + _remote_forward, + args, + kwargs, + ) + return ret_fut.wait() + + +_generated_methods = [ + forward_async, + forward, +] + + + + +def _remote_forward( + module_rref: RRef[module_interface_cls], device: str, is_device_map_set: bool, *args, **kwargs): + module = module_rref.local_value() + device = torch.device(device) + + if device.type != "cuda": + return module.forward(*args, **kwargs) + + # If the module is on a cuda device, + # move any CPU tensor in args or kwargs to the same cuda device. + # Since torch script does not support generator expression, + # have to use concatenation instead of + # ``tuple(i.to(device) if isinstance(i, Tensor) else i for i in *args)``. + args = (*args,) + out_args: Tuple[()] = () + for arg in args: + arg = (arg.to(device),) if isinstance(arg, Tensor) else (arg,) + out_args = out_args + arg + + kwargs = {**kwargs} + for k, v in kwargs.items(): + if isinstance(v, Tensor): + kwargs[k] = kwargs[k].to(device) + + if is_device_map_set: + return module.forward(*out_args, **kwargs) + + # If the device map is empty, then only CPU tensors are allowed to send over wire, + # so have to move any GPU tensor to CPU in the output. + # Since torch script does not support generator expression, + # have to use concatenation instead of + # ``tuple(i.cpu() if isinstance(i, Tensor) else i for i in module.forward(*out_args, **kwargs))``. + ret: Tuple[()] = () + for i in module.forward(*out_args, **kwargs): + i = (i.cpu(),) if isinstance(i, Tensor) else (i,) + ret = ret + i + return ret diff --git a/tmpqpwscwjm/__pycache__/_remote_module_non_scriptable.cpython-312.pyc b/tmpqpwscwjm/__pycache__/_remote_module_non_scriptable.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..5ac75c8b2bb560f1364841ef7f448e39584f654a Binary files /dev/null and b/tmpqpwscwjm/__pycache__/_remote_module_non_scriptable.cpython-312.pyc differ diff --git a/tmpqpwscwjm/_remote_module_non_scriptable.py b/tmpqpwscwjm/_remote_module_non_scriptable.py new file mode 100644 index 0000000000000000000000000000000000000000..796e049370abda22e60acbeda5d65a4be33e1a87 --- /dev/null +++ b/tmpqpwscwjm/_remote_module_non_scriptable.py @@ -0,0 +1,81 @@ +from typing import * + +import torch +import torch.distributed.rpc as rpc +from torch import Tensor +from torch._jit_internal import Future +from torch.distributed.rpc import RRef +from typing import Tuple # pyre-ignore: unused import + + +module_interface_cls = None + + +def forward_async(self, *args, **kwargs): + args = (self.module_rref, self.device, self.is_device_map_set, *args) + kwargs = {**kwargs} + return rpc.rpc_async( + self.module_rref.owner(), + _remote_forward, + args, + kwargs, + ) + + +def forward(self, *args, **kwargs): + args = (self.module_rref, self.device, self.is_device_map_set, *args) + kwargs = {**kwargs} + ret_fut = rpc.rpc_async( + self.module_rref.owner(), + _remote_forward, + args, + kwargs, + ) + return ret_fut.wait() + + +_generated_methods = [ + forward_async, + forward, +] + + + + +def _remote_forward( + module_rref: RRef[module_interface_cls], device: str, is_device_map_set: bool, *args, **kwargs): + module = module_rref.local_value() + device = torch.device(device) + + if device.type != "cuda": + return module.forward(*args, **kwargs) + + # If the module is on a cuda device, + # move any CPU tensor in args or kwargs to the same cuda device. + # Since torch script does not support generator expression, + # have to use concatenation instead of + # ``tuple(i.to(device) if isinstance(i, Tensor) else i for i in *args)``. + args = (*args,) + out_args: Tuple[()] = () + for arg in args: + arg = (arg.to(device),) if isinstance(arg, Tensor) else (arg,) + out_args = out_args + arg + + kwargs = {**kwargs} + for k, v in kwargs.items(): + if isinstance(v, Tensor): + kwargs[k] = kwargs[k].to(device) + + if is_device_map_set: + return module.forward(*out_args, **kwargs) + + # If the device map is empty, then only CPU tensors are allowed to send over wire, + # so have to move any GPU tensor to CPU in the output. + # Since torch script does not support generator expression, + # have to use concatenation instead of + # ``tuple(i.cpu() if isinstance(i, Tensor) else i for i in module.forward(*out_args, **kwargs))``. + ret: Tuple[()] = () + for i in module.forward(*out_args, **kwargs): + i = (i.cpu(),) if isinstance(i, Tensor) else (i,) + ret = ret + i + return ret diff --git a/tmptu2tppnl/__pycache__/_remote_module_non_scriptable.cpython-312.pyc b/tmptu2tppnl/__pycache__/_remote_module_non_scriptable.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..2f78b92541a6774e04d65a5ea14f7b1c25f55a56 Binary files /dev/null and b/tmptu2tppnl/__pycache__/_remote_module_non_scriptable.cpython-312.pyc differ diff --git a/tmptu2tppnl/_remote_module_non_scriptable.py b/tmptu2tppnl/_remote_module_non_scriptable.py new file mode 100644 index 0000000000000000000000000000000000000000..796e049370abda22e60acbeda5d65a4be33e1a87 --- /dev/null +++ b/tmptu2tppnl/_remote_module_non_scriptable.py @@ -0,0 +1,81 @@ +from typing import * + +import torch +import torch.distributed.rpc as rpc +from torch import Tensor +from torch._jit_internal import Future +from torch.distributed.rpc import RRef +from typing import Tuple # pyre-ignore: unused import + + +module_interface_cls = None + + +def forward_async(self, *args, **kwargs): + args = (self.module_rref, self.device, self.is_device_map_set, *args) + kwargs = {**kwargs} + return rpc.rpc_async( + self.module_rref.owner(), + _remote_forward, + args, + kwargs, + ) + + +def forward(self, *args, **kwargs): + args = (self.module_rref, self.device, self.is_device_map_set, *args) + kwargs = {**kwargs} + ret_fut = rpc.rpc_async( + self.module_rref.owner(), + _remote_forward, + args, + kwargs, + ) + return ret_fut.wait() + + +_generated_methods = [ + forward_async, + forward, +] + + + + +def _remote_forward( + module_rref: RRef[module_interface_cls], device: str, is_device_map_set: bool, *args, **kwargs): + module = module_rref.local_value() + device = torch.device(device) + + if device.type != "cuda": + return module.forward(*args, **kwargs) + + # If the module is on a cuda device, + # move any CPU tensor in args or kwargs to the same cuda device. + # Since torch script does not support generator expression, + # have to use concatenation instead of + # ``tuple(i.to(device) if isinstance(i, Tensor) else i for i in *args)``. + args = (*args,) + out_args: Tuple[()] = () + for arg in args: + arg = (arg.to(device),) if isinstance(arg, Tensor) else (arg,) + out_args = out_args + arg + + kwargs = {**kwargs} + for k, v in kwargs.items(): + if isinstance(v, Tensor): + kwargs[k] = kwargs[k].to(device) + + if is_device_map_set: + return module.forward(*out_args, **kwargs) + + # If the device map is empty, then only CPU tensors are allowed to send over wire, + # so have to move any GPU tensor to CPU in the output. + # Since torch script does not support generator expression, + # have to use concatenation instead of + # ``tuple(i.cpu() if isinstance(i, Tensor) else i for i in module.forward(*out_args, **kwargs))``. + ret: Tuple[()] = () + for i in module.forward(*out_args, **kwargs): + i = (i.cpu(),) if isinstance(i, Tensor) else (i,) + ret = ret + i + return ret diff --git a/tmpvkx5qgfg/__pycache__/_remote_module_non_scriptable.cpython-312.pyc b/tmpvkx5qgfg/__pycache__/_remote_module_non_scriptable.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..8b5cbde8fc3fe15995e0141f2c25efe67b15ba09 Binary files /dev/null and b/tmpvkx5qgfg/__pycache__/_remote_module_non_scriptable.cpython-312.pyc differ diff --git a/tmpvkx5qgfg/_remote_module_non_scriptable.py b/tmpvkx5qgfg/_remote_module_non_scriptable.py new file mode 100644 index 0000000000000000000000000000000000000000..796e049370abda22e60acbeda5d65a4be33e1a87 --- /dev/null +++ b/tmpvkx5qgfg/_remote_module_non_scriptable.py @@ -0,0 +1,81 @@ +from typing import * + +import torch +import torch.distributed.rpc as rpc +from torch import Tensor +from torch._jit_internal import Future +from torch.distributed.rpc import RRef +from typing import Tuple # pyre-ignore: unused import + + +module_interface_cls = None + + +def forward_async(self, *args, **kwargs): + args = (self.module_rref, self.device, self.is_device_map_set, *args) + kwargs = {**kwargs} + return rpc.rpc_async( + self.module_rref.owner(), + _remote_forward, + args, + kwargs, + ) + + +def forward(self, *args, **kwargs): + args = (self.module_rref, self.device, self.is_device_map_set, *args) + kwargs = {**kwargs} + ret_fut = rpc.rpc_async( + self.module_rref.owner(), + _remote_forward, + args, + kwargs, + ) + return ret_fut.wait() + + +_generated_methods = [ + forward_async, + forward, +] + + + + +def _remote_forward( + module_rref: RRef[module_interface_cls], device: str, is_device_map_set: bool, *args, **kwargs): + module = module_rref.local_value() + device = torch.device(device) + + if device.type != "cuda": + return module.forward(*args, **kwargs) + + # If the module is on a cuda device, + # move any CPU tensor in args or kwargs to the same cuda device. + # Since torch script does not support generator expression, + # have to use concatenation instead of + # ``tuple(i.to(device) if isinstance(i, Tensor) else i for i in *args)``. + args = (*args,) + out_args: Tuple[()] = () + for arg in args: + arg = (arg.to(device),) if isinstance(arg, Tensor) else (arg,) + out_args = out_args + arg + + kwargs = {**kwargs} + for k, v in kwargs.items(): + if isinstance(v, Tensor): + kwargs[k] = kwargs[k].to(device) + + if is_device_map_set: + return module.forward(*out_args, **kwargs) + + # If the device map is empty, then only CPU tensors are allowed to send over wire, + # so have to move any GPU tensor to CPU in the output. + # Since torch script does not support generator expression, + # have to use concatenation instead of + # ``tuple(i.cpu() if isinstance(i, Tensor) else i for i in module.forward(*out_args, **kwargs))``. + ret: Tuple[()] = () + for i in module.forward(*out_args, **kwargs): + i = (i.cpu(),) if isinstance(i, Tensor) else (i,) + ret = ret + i + return ret diff --git a/tmpxk_w0eue/__pycache__/_remote_module_non_scriptable.cpython-312.pyc b/tmpxk_w0eue/__pycache__/_remote_module_non_scriptable.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..a44067790d1fed7cfbfc560bc5400952a834a28e Binary files /dev/null and b/tmpxk_w0eue/__pycache__/_remote_module_non_scriptable.cpython-312.pyc differ diff --git a/tmpxk_w0eue/_remote_module_non_scriptable.py b/tmpxk_w0eue/_remote_module_non_scriptable.py new file mode 100644 index 0000000000000000000000000000000000000000..796e049370abda22e60acbeda5d65a4be33e1a87 --- /dev/null +++ b/tmpxk_w0eue/_remote_module_non_scriptable.py @@ -0,0 +1,81 @@ +from typing import * + +import torch +import torch.distributed.rpc as rpc +from torch import Tensor +from torch._jit_internal import Future +from torch.distributed.rpc import RRef +from typing import Tuple # pyre-ignore: unused import + + +module_interface_cls = None + + +def forward_async(self, *args, **kwargs): + args = (self.module_rref, self.device, self.is_device_map_set, *args) + kwargs = {**kwargs} + return rpc.rpc_async( + self.module_rref.owner(), + _remote_forward, + args, + kwargs, + ) + + +def forward(self, *args, **kwargs): + args = (self.module_rref, self.device, self.is_device_map_set, *args) + kwargs = {**kwargs} + ret_fut = rpc.rpc_async( + self.module_rref.owner(), + _remote_forward, + args, + kwargs, + ) + return ret_fut.wait() + + +_generated_methods = [ + forward_async, + forward, +] + + + + +def _remote_forward( + module_rref: RRef[module_interface_cls], device: str, is_device_map_set: bool, *args, **kwargs): + module = module_rref.local_value() + device = torch.device(device) + + if device.type != "cuda": + return module.forward(*args, **kwargs) + + # If the module is on a cuda device, + # move any CPU tensor in args or kwargs to the same cuda device. + # Since torch script does not support generator expression, + # have to use concatenation instead of + # ``tuple(i.to(device) if isinstance(i, Tensor) else i for i in *args)``. + args = (*args,) + out_args: Tuple[()] = () + for arg in args: + arg = (arg.to(device),) if isinstance(arg, Tensor) else (arg,) + out_args = out_args + arg + + kwargs = {**kwargs} + for k, v in kwargs.items(): + if isinstance(v, Tensor): + kwargs[k] = kwargs[k].to(device) + + if is_device_map_set: + return module.forward(*out_args, **kwargs) + + # If the device map is empty, then only CPU tensors are allowed to send over wire, + # so have to move any GPU tensor to CPU in the output. + # Since torch script does not support generator expression, + # have to use concatenation instead of + # ``tuple(i.cpu() if isinstance(i, Tensor) else i for i in module.forward(*out_args, **kwargs))``. + ret: Tuple[()] = () + for i in module.forward(*out_args, **kwargs): + i = (i.cpu(),) if isinstance(i, Tensor) else (i,) + ret = ret + i + return ret diff --git a/tokenizer.json b/tokenizer.json new file mode 100644 index 0000000000000000000000000000000000000000..570d3a247dd74fe76cc5340c973cde2c1c6601b9 --- /dev/null +++ b/tokenizer.json @@ -0,0 +1,133 @@ +{ + "version": "1.0", + "truncation": null, + "padding": null, + "added_tokens": [ + { + "id": 0, + "content": "", + "single_word": false, + "lstrip": false, + "rstrip": false, + "normalized": false, + "special": true + }, + { + "id": 1, + "content": "", + "single_word": false, + "lstrip": false, + "rstrip": false, + "normalized": false, + "special": true + }, + { + "id": 2, + "content": "", + "single_word": false, + "lstrip": false, + "rstrip": false, + "normalized": false, + "special": true + }, + { + "id": 3, + "content": "", + "single_word": false, + "lstrip": false, + "rstrip": false, + "normalized": false, + "special": true + }, + { + "id": 33, + "content": "<+>", + "single_word": false, + "lstrip": false, + "rstrip": false, + "normalized": false, + "special": true + }, + { + "id": 34, + "content": "<->", + "single_word": false, + "lstrip": false, + "rstrip": false, + "normalized": false, + "special": true + }, + { + "id": 35, + "content": "", + "single_word": false, + "lstrip": false, + "rstrip": false, + "normalized": false, + "special": true + }, + { + "id": 36, + "content": "", + "single_word": false, + "lstrip": false, + "rstrip": false, + "normalized": false, + "special": true + } + ], + "normalizer": null, + "pre_tokenizer": null, + "post_processor": null, + "decoder": null, + "model": { + "type": "BPE", + "dropout": null, + "unk_token": "", + "continuing_subword_prefix": null, + "end_of_word_suffix": null, + "fuse_unk": false, + "byte_fallback": false, + "ignore_merges": false, + "vocab": { + "": 0, + "": 1, + "": 2, + "": 3, + "L": 4, + "A": 5, + "G": 6, + "V": 7, + "S": 8, + "E": 9, + "R": 10, + "T": 11, + "I": 12, + "D": 13, + "P": 14, + "K": 15, + "Q": 16, + "N": 17, + "F": 18, + "Y": 19, + "M": 20, + "H": 21, + "W": 22, + "C": 23, + "X": 24, + "B": 25, + "U": 26, + "Z": 27, + "O": 28, + "a": 29, + "t": 30, + "c": 31, + "g": 32, + "<+>": 33, + "<->": 34, + "": 35, + "": 36 + }, + "merges": [] + } +} \ No newline at end of file diff --git a/tokenizer_config.json b/tokenizer_config.json new file mode 100644 index 0000000000000000000000000000000000000000..4d901f7c793d57ccbdd6b67298f86dcf3940579a --- /dev/null +++ b/tokenizer_config.json @@ -0,0 +1,84 @@ +{ + "added_tokens_decoder": { + "0": { + "content": "", + "lstrip": false, + "normalized": false, + "rstrip": false, + "single_word": false, + "special": true + }, + "1": { + "content": "", + "lstrip": false, + "normalized": false, + "rstrip": false, + "single_word": false, + "special": true + }, + "2": { + "content": "", + "lstrip": false, + "normalized": false, + "rstrip": false, + "single_word": false, + "special": true + }, + "3": { + "content": "", + "lstrip": false, + "normalized": false, + "rstrip": false, + "single_word": false, + "special": true + }, + "33": { + "content": "<+>", + "lstrip": false, + "normalized": false, + "rstrip": false, + "single_word": false, + "special": true + }, + "34": { + "content": "<->", + "lstrip": false, + "normalized": false, + "rstrip": false, + "single_word": false, + "special": true + }, + "35": { + "content": "", + "lstrip": false, + "normalized": false, + "rstrip": false, + "single_word": false, + "special": true + }, + "36": { + "content": "", + "lstrip": false, + "normalized": false, + "rstrip": false, + "single_word": false, + "special": true + } + }, + "auto_map": { + "AutoTokenizer": [ + "glm_tokenizer.gLM2Tokenizer", + null + ] + }, + "clean_up_tokenization_spaces": true, + "cls_token": "", + "eos_token": "", + "extra_special_tokens": {}, + "mask_token": "", + "model_max_length": 1000000000000000019884624838656, + "pad_token": "", + "sep_token": "", + "tokenizer_class": "gLM2Tokenizer", + "unk_token": "" +}