{"metadata":{"kernelspec":{"language":"python","display_name":"Python 3","name":"python3"},"language_info":{"name":"python","version":"3.10.12","mimetype":"text/x-python","codemirror_mode":{"name":"ipython","version":3},"pygments_lexer":"ipython3","nbconvert_exporter":"python","file_extension":".py"},"kaggle":{"accelerator":"gpu","dataSources":[{"sourceId":87793,"databundleVersionId":11403143,"isSourceIdPinned":false,"sourceType":"competition"},{"sourceId":11029791,"sourceType":"datasetVersion","datasetId":6869176},{"sourceId":11031805,"sourceType":"datasetVersion","datasetId":6870779},{"sourceId":11037781,"sourceType":"datasetVersion","datasetId":6875047}],"dockerImageVersionId":30918,"isInternetEnabled":false,"language":"python","sourceType":"notebook","isGpuEnabled":true}},"nbformat_minor":4,"nbformat":4,"cells":[{"cell_type":"code","source":"!pip install /kaggle/input/biopython/biopython-1.85-cp310-cp310-manylinux_2_17_x86_64.manylinux2014_x86_64.whl","metadata":{"_uuid":"8f2839f25d086af736a60e9eeb907d3b93b6e0e5","_cell_guid":"b1076dfc-b9ad-4769-8c92-a6c4dae69d19","trusted":true,"execution":{"iopub.status.busy":"2025-03-15T07:59:35.379396Z","iopub.execute_input":"2025-03-15T07:59:35.379686Z","iopub.status.idle":"2025-03-15T07:59:40.745376Z","shell.execute_reply.started":"2025-03-15T07:59:35.379662Z","shell.execute_reply":"2025-03-15T07:59:40.744552Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"import pandas as pd\nimport numpy as np\nimport torch\nimport os\nfrom tqdm import tqdm\nfrom torch.utils.data import Dataset, DataLoader\nimport random\nimport torch.nn as nn\nimport torch.nn.functional as F\n\n# 재현성을 위한 시드 설정\ndef set_seed(seed=42):\n    random.seed(seed)\n    np.random.seed(seed)\n    torch.manual_seed(seed)\n    torch.cuda.manual_seed(seed)\n    torch.backends.cudnn.deterministic = True\n    torch.backends.cudnn.benchmark = False\n\n# 설정\nconfig = {\n    \"max_len\": 384,\n    \"msa_max_sequences\": 32,\n    \"msa_feat_dim\": 128,\n    \"num_self_attn_layers\": 4,\n    \"num_cross_attn_layers\": 4,\n    \"num_structure_module_layers\": 8,\n    \"n_heads\": 8,\n    \"dropout\": 0.25,\n    \"num_ensemble\": 5,  # 5개의 앙상블 예측\n}","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-03-15T07:59:40.747121Z","iopub.execute_input":"2025-03-15T07:59:40.747351Z","iopub.status.idle":"2025-03-15T07:59:44.249830Z","shell.execute_reply.started":"2025-03-15T07:59:40.747330Z","shell.execute_reply":"2025-03-15T07:59:44.249071Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# BioPython을 사용한 FASTA/MSA 파일 로드 함수\ndef load_msa_with_biopython(fasta_file, convert_to_rna=False):\n    \"\"\"\n    BioPython의 SeqIO를 사용하여 FASTA 형식의 MSA 파일에서 시퀀스를 로드합니다.\n    \"\"\"\n    try:\n        from Bio import SeqIO\n        sequences = []\n        \n        # FASTA 파일에서 레코드 읽기\n        for record in SeqIO.parse(fasta_file, \"fasta\"):\n            # 시퀀스를 문자열로 변환\n            seq_str = str(record.seq).upper()\n            \n            # DNA를 RNA로 변환 (필요한 경우)\n            if convert_to_rna:\n                seq_str = seq_str.replace('T', 'U')\n            \n            sequences.append(seq_str)\n        \n        # 시퀀스가 없는 경우 기본값 반환\n        if not sequences:\n            return [\"\"]\n        \n        return sequences\n    \n    except ImportError:\n        print(\"BioPython이 설치되어 있지 않습니다. 기본 파서를 사용합니다.\")\n        return load_a3m_fallback(fasta_file, convert_to_rna)\n    except Exception as e:\n        print(f\"FASTA 파일 로드 중 오류 발생: {str(e)}\")\n        return [\"\"]\n\n# BioPython 없을 경우를 위한 대체 파서\ndef load_a3m_fallback(fasta_file, convert_to_rna=False):\n    \"\"\"기본 파서를 사용하여 FASTA 파일을 로드합니다.\"\"\"\n    sequences = []\n    current_seq = \"\"\n    \n    with open(fasta_file, \"r\") as f:\n        for line in f:\n            line = line.strip()\n            if line.startswith(\">\"):\n                if current_seq:\n                    sequences.append(current_seq)\n                    current_seq = \"\"\n            else:\n                # 대문자로 정규화\n                clean_line = line.upper()\n                # DNA를 RNA로 변환 (필요한 경우)\n                if convert_to_rna:\n                    clean_line = clean_line.replace('T', 'U')\n                \n                current_seq += clean_line\n    \n    if current_seq:\n        sequences.append(current_seq)\n    \n    # 시퀀스가 없는 경우를 대비\n    if not sequences:\n        return [\"\"]\n    \n    return sequences\n\n# MSA 데이터를 처리하는 함수\ndef process_msa(msa_sequences, max_seq=32):\n    \"\"\"MSA 시퀀스를 처리하여 숫자 인코딩 텐서로 변환합니다.\"\"\"\n    # 시퀀스가 없거나 빈 문자열인 경우 처리\n    if not msa_sequences or not msa_sequences[0]:\n        # 기본값으로 1x1 배열 반환\n        return np.zeros((1, 1), dtype=np.int64)\n    \n    # 첫 번째 시퀀스는 쿼리 시퀀스\n    query_seq = msa_sequences[0]\n    \n    # 최대 시퀀스 수 제한\n    msa_sequences = msa_sequences[:max_seq]\n    \n    # 시퀀스 길이 확인\n    seq_len = len(query_seq)\n    \n    # 시퀀스가 비어있으면 기본값 반환\n    if seq_len == 0:\n        return np.zeros((len(msa_sequences), 1), dtype=np.int64)\n    \n    # 원-핫 인코딩을 위한 사전 (A, C, G, U, - 갭, 기타 문자)\n    vocab = {'A': 0, 'C': 1, 'G': 2, 'U': 3, 'T': 4, '-': 5, '.': 5}  # Add 'T' as token 4\n    \n    # MSA 시퀀스를 숫자로 변환\n    msa_encoded = np.zeros((len(msa_sequences), seq_len), dtype=np.int64)\n    \n    # 각 시퀀스 처리\n    for i, seq in enumerate(msa_sequences):\n        # 현재 시퀀스가 쿼리 시퀀스보다 짧으면 패딩\n        if len(seq) < seq_len:\n            seq = seq + '-' * (seq_len - len(seq))\n        \n        # 시퀀스가 쿼리보다 길면 잘라내기\n        seq = seq[:seq_len]\n        \n        # 문자별 처리\n        for j, nt in enumerate(seq):\n            nt_upper = nt.upper()\n            \n            if nt_upper in vocab:\n                msa_encoded[i, j] = vocab[nt_upper]\n            else:\n                # 모르는 문자는 갭으로 처리\n                msa_encoded[i, j] = vocab['-']\n    \n    return msa_encoded\n\n# 테스트 데이터셋 클래스\nclass RNA3D_MSA_TestDataset(Dataset):\n    def __init__(self, test_data, msa_dir, config):\n        self.test_data = test_data\n        self.msa_dir = msa_dir\n        self.config = config\n        # Line ~166: Change the tokens dictionary \n        self.tokens = {'A': 0, 'C': 1, 'G': 2, 'U': 3, 'T': 4}  # Add 'T'\n    \n    def __len__(self):\n        return len(self.test_data)\n    \n    def __getitem__(self, idx):\n        target_id = self.test_data.loc[idx, 'target_id']\n        sequence = self.test_data.loc[idx, 'sequence']\n        \n        # 시퀀스를 숫자로 변환\n        seq_encoded = [self.tokens.get(nt, 0) for nt in sequence]\n        seq_encoded = torch.tensor(np.array(seq_encoded))\n        \n        # MSA 파일 경로\n        msa_file = os.path.join(self.msa_dir, f\"{target_id}.MSA.fasta\")\n        \n        # MSA 데이터 로드\n        if os.path.exists(msa_file):\n            try:\n                msa_sequences = load_msa_with_biopython(msa_file)\n                msa_encoded = process_msa(msa_sequences, self.config['msa_max_sequences'])\n                msa_encoded = torch.tensor(msa_encoded)\n            except Exception as e:\n                print(f\"MSA 로드 실패: {target_id}, 오류: {str(e)}\")\n                # 기본값: 원본 시퀀스만 포함\n                msa_encoded = torch.tensor([[self.tokens.get(nt, 0) for nt in sequence]])\n        else:\n            # MSA 파일이 없는 경우, 원래 시퀀스만 포함\n            msa_encoded = torch.tensor([[self.tokens.get(nt, 0) for nt in sequence]])\n        \n        return {\n            'target_id': target_id,\n            'sequence': seq_encoded,\n            'msa': msa_encoded,\n            'raw_sequence': sequence\n        }\n\n# 모델 정의 (학습 코드와 동일)\nclass MSAAttention(nn.Module):\n    def __init__(self, d_model, n_heads, dropout=0.1):\n        super(MSAAttention, self).__init__()\n        self.d_model = d_model\n        self.n_heads = n_heads\n        self.head_dim = d_model // n_heads\n        \n        self.query = nn.Linear(d_model, d_model)\n        self.key = nn.Linear(d_model, d_model)\n        self.value = nn.Linear(d_model, d_model)\n        self.out = nn.Linear(d_model, d_model)\n        \n        self.dropout = nn.Dropout(dropout)\n    \n    def forward(self, x, mask=None):\n        batch_size, num_seqs, seq_len, _ = x.shape\n        \n        q = self.query(x).view(batch_size, num_seqs, seq_len, self.n_heads, self.head_dim)\n        k = self.key(x).view(batch_size, num_seqs, seq_len, self.n_heads, self.head_dim)\n        v = self.value(x).view(batch_size, num_seqs, seq_len, self.n_heads, self.head_dim)\n        \n        # 차원 변경: batch_size, n_heads, num_seqs, seq_len, head_dim\n        q = q.permute(0, 3, 1, 2, 4)\n        k = k.permute(0, 3, 1, 2, 4)\n        v = v.permute(0, 3, 1, 2, 4)\n        \n        # 어텐션 계산\n        scores = torch.matmul(q, k.transpose(-2, -1)) / (self.head_dim ** 0.5)\n        \n        if mask is not None:\n            scores = scores.masked_fill(mask == 0, -1e9)\n        \n        attn_weights = F.softmax(scores, dim=-1)\n        attn_weights = self.dropout(attn_weights)\n        \n        output = torch.matmul(attn_weights, v)\n        \n        # 차원 되돌리기\n        output = output.permute(0, 2, 3, 1, 4).contiguous()\n        output = output.view(batch_size, num_seqs, seq_len, self.d_model)\n        \n        return self.out(output)\n\nclass RowAttention(nn.Module):\n    \"\"\"시퀀스 위치 간의 어텐션 (각 MSA 시퀀스 내에서)\"\"\"\n    def __init__(self, d_model, n_heads, dropout=0.1):\n        super(RowAttention, self).__init__()\n        self.attention = MSAAttention(d_model, n_heads, dropout)\n    \n    def forward(self, x, mask=None):\n        return self.attention(x, mask)\n\nclass ColumnAttention(nn.Module):\n    \"\"\"MSA 시퀀스 간의 어텐션 (각 위치에서)\"\"\"\n    def __init__(self, d_model, n_heads, dropout=0.1):\n        super(ColumnAttention, self).__init__()\n        self.attention = MSAAttention(d_model, n_heads, dropout)\n    \n    def forward(self, x, mask=None):\n        batch_size, num_seqs, seq_len, d_model = x.shape\n        \n        # 차원 변경: (batch_size, seq_len, num_seqs, d_model)\n        x = x.permute(0, 2, 1, 3)\n        \n        if mask is not None:\n            mask = mask.permute(0, 2, 1, 3)\n        \n        x = self.attention(x, mask)\n        \n        # 차원 되돌리기\n        x = x.permute(0, 2, 1, 3)\n        \n        return x\n\nclass FeedForward(nn.Module):\n    def __init__(self, d_model, d_ff=2048, dropout=0.1):\n        super(FeedForward, self).__init__()\n        self.linear1 = nn.Linear(d_model, d_ff)\n        self.dropout = nn.Dropout(dropout)\n        self.linear2 = nn.Linear(d_ff, d_model)\n    \n    def forward(self, x):\n        return self.linear2(self.dropout(F.relu(self.linear1(x))))\n\nclass LayerNorm(nn.Module):\n    def __init__(self, d_model, eps=1e-6):\n        super(LayerNorm, self).__init__()\n        self.gamma = nn.Parameter(torch.ones(d_model))\n        self.beta = nn.Parameter(torch.zeros(d_model))\n        self.eps = eps\n    \n    def forward(self, x):\n        mean = x.mean(-1, keepdim=True)\n        std = x.std(-1, keepdim=True)\n        return self.gamma * (x - mean) / (std + self.eps) + self.beta\n\nclass MSATransformerLayer(nn.Module):\n    def __init__(self, d_model, n_heads, d_ff=2048, dropout=0.1):\n        super(MSATransformerLayer, self).__init__()\n        self.row_attn = RowAttention(d_model, n_heads, dropout)\n        self.col_attn = ColumnAttention(d_model, n_heads, dropout)\n        self.feed_forward = FeedForward(d_model, d_ff, dropout)\n        \n        self.norm1 = LayerNorm(d_model)\n        self.norm2 = LayerNorm(d_model)\n        self.norm3 = LayerNorm(d_model)\n        \n        self.dropout1 = nn.Dropout(dropout)\n        self.dropout2 = nn.Dropout(dropout)\n        self.dropout3 = nn.Dropout(dropout)\n    \n    def forward(self, x, mask=None):\n        # 행 어텐션 (시퀀스 위치 간)\n        row_attn_output = x + self.dropout1(self.row_attn(self.norm1(x), mask))\n        \n        # 열 어텐션 (MSA 시퀀스 간)\n        col_attn_output = row_attn_output + self.dropout2(self.col_attn(self.norm2(row_attn_output), mask))\n        \n        # 피드포워드\n        output = col_attn_output + self.dropout3(self.feed_forward(self.norm3(col_attn_output)))\n        \n        return output\n\nclass MSATransformer(nn.Module):\n    def __init__(self, d_model, n_heads, num_layers, d_ff=2048, dropout=0.1):\n        super(MSATransformer, self).__init__()\n        self.layers = nn.ModuleList([\n            MSATransformerLayer(d_model, n_heads, d_ff, dropout)\n            for _ in range(num_layers)\n        ])\n        self.norm = LayerNorm(d_model)\n    \n    def forward(self, x, mask=None):\n        for layer in self.layers:\n            x = layer(x, mask)\n        return self.norm(x)\n\nclass CrossAttention(nn.Module):\n    def __init__(self, d_model, n_heads, dropout=0.1):\n        super(CrossAttention, self).__init__()\n        self.d_model = d_model\n        self.n_heads = n_heads\n        self.head_dim = d_model // n_heads\n        \n        self.query = nn.Linear(d_model, d_model)\n        self.key = nn.Linear(d_model, d_model)\n        self.value = nn.Linear(d_model, d_model)\n        self.out = nn.Linear(d_model, d_model)\n        \n        self.dropout = nn.Dropout(dropout)\n    \n    def forward(self, query, key_value, mask=None):\n        batch_size = query.shape[0]\n        \n        q = self.query(query)\n        k = self.key(key_value)\n        v = self.value(key_value)\n        \n        # 어텐션 차원 변경\n        q = q.view(batch_size, -1, self.n_heads, self.head_dim).transpose(1, 2)\n        k = k.view(batch_size, -1, self.n_heads, self.head_dim).transpose(1, 2)\n        v = v.view(batch_size, -1, self.n_heads, self.head_dim).transpose(1, 2)\n        \n        # 어텐션 계산\n        scores = torch.matmul(q, k.transpose(-2, -1)) / (self.head_dim ** 0.5)\n        \n        if mask is not None:\n            scores = scores.masked_fill(mask == 0, -1e9)\n        \n        attn_weights = F.softmax(scores, dim=-1)\n        attn_weights = self.dropout(attn_weights)\n        \n        output = torch.matmul(attn_weights, v)\n        \n        # 차원 되돌리기\n        output = output.transpose(1, 2).contiguous()\n        output = output.view(batch_size, -1, self.d_model)\n        \n        return self.out(output)\n\nclass CrossTransformerLayer(nn.Module):\n    def __init__(self, d_model, n_heads, d_ff=2048, dropout=0.1):\n        super(CrossTransformerLayer, self).__init__()\n        self.cross_attn = CrossAttention(d_model, n_heads, dropout)\n        self.feed_forward = FeedForward(d_model, d_ff, dropout)\n        \n        self.norm1 = LayerNorm(d_model)\n        self.norm2 = LayerNorm(d_model)\n        self.norm3 = LayerNorm(d_model)\n        \n        self.dropout1 = nn.Dropout(dropout)\n        self.dropout2 = nn.Dropout(dropout)\n    \n    def forward(self, x, msa_repr, mask=None):\n        # 크로스 어텐션\n        cross_attn_output = x + self.dropout1(self.cross_attn(self.norm1(x), self.norm2(msa_repr), mask))\n        \n        # 피드포워드\n        output = cross_attn_output + self.dropout2(self.feed_forward(self.norm3(cross_attn_output)))\n        \n        return output\n\nclass CrossTransformer(nn.Module):\n    def __init__(self, d_model, n_heads, num_layers, d_ff=2048, dropout=0.1):\n        super(CrossTransformer, self).__init__()\n        self.layers = nn.ModuleList([\n            CrossTransformerLayer(d_model, n_heads, d_ff, dropout)\n            for _ in range(num_layers)\n        ])\n        self.norm = LayerNorm(d_model)\n    \n    def forward(self, x, msa_repr, mask=None):\n        for layer in self.layers:\n            x = layer(x, msa_repr, mask)\n        return self.norm(x)\n\nclass StructureModule(nn.Module):\n    def __init__(self, d_model, n_heads, num_layers, d_ff=2048, dropout=0.1):\n        super(StructureModule, self).__init__()\n        self.layers = nn.ModuleList([\n            MSATransformerLayer(d_model, n_heads, d_ff, dropout)\n            for _ in range(num_layers)\n        ])\n        self.norm = LayerNorm(d_model)\n        self.xyz_predictor = nn.Linear(d_model, 3)\n    \n    def forward(self, x, mask=None):\n        for layer in self.layers:\n            x = layer(x, mask)\n        x = self.norm(x)\n        xyz = self.xyz_predictor(x)\n        return xyz\n\nclass RNA_MSA_Folding(nn.Module):\n    def __init__(self, config):\n        super(RNA_MSA_Folding, self).__init__()\n        self.config = config\n        \n        # 임베딩 레이어\n        self.seq_embedding = nn.Embedding(6, config['msa_feat_dim'])  # Change from 5 to 6 to include 'T'\n\n        \n        # 위치 인코딩\n        self.pos_embedding = nn.Parameter(torch.zeros(1, 1, config['max_len'], config['msa_feat_dim']))\n        \n        # MSA Transformer\n        self.msa_transformer = MSATransformer(\n            d_model=config['msa_feat_dim'],\n            n_heads=config['n_heads'],\n            num_layers=config['num_self_attn_layers'],\n            dropout=config['dropout']\n        )\n        \n        # 단일 시퀀스 임베딩을 위한 추가 레이어\n        self.single_seq_embedding = nn.Linear(config['msa_feat_dim'], config['msa_feat_dim'])\n        \n        # Cross Attention\n        self.cross_transformer = CrossTransformer(\n            d_model=config['msa_feat_dim'],\n            n_heads=config['n_heads'],\n            num_layers=config['num_cross_attn_layers'],\n            dropout=config['dropout']\n        )\n        \n        # 구조 모듈\n        self.structure_module = StructureModule(\n            d_model=config['msa_feat_dim'],\n            n_heads=config['n_heads'],\n            num_layers=config['num_structure_module_layers'],\n            dropout=config['dropout']\n        )\n    \n    def forward(self, seq, msa):\n        # 입력 형태 확인 및 조정\n        if len(msa.shape) != 3:\n            # 최소 3차원 텐서로 조정\n            if len(msa.shape) == 2:\n                msa = msa.unsqueeze(0)\n            elif len(msa.shape) == 1:\n                msa = msa.unsqueeze(0).unsqueeze(0)\n        \n        batch_size, num_seqs, seq_len = msa.shape\n        \n        # MSA 임베딩\n        msa_emb = self.seq_embedding(msa)  # [batch_size, num_seqs, seq_len, d_model]\n        \n        # 위치 인코딩 추가 (너무 긴 시퀀스 처리)\n        if seq_len <= self.pos_embedding.shape[2]:\n            pos_emb = self.pos_embedding[:, :, :seq_len, :]\n        else:\n            # 위치 인코딩 확장\n            existing_pos = self.pos_embedding.squeeze(0).squeeze(0)  # [max_len, d_model]\n            extended_pos = torch.zeros(seq_len, self.config['msa_feat_dim'], device=msa.device)\n            extended_pos[:existing_pos.shape[0], :] = existing_pos\n            # 나머지 부분은 마지막 위치 인코딩 복제\n            if existing_pos.shape[0] > 0:\n                extended_pos[existing_pos.shape[0]:, :] = existing_pos[-1, :]\n            pos_emb = extended_pos.unsqueeze(0).unsqueeze(0)\n            \n        msa_emb = msa_emb + pos_emb\n        \n        # MSA 처리\n        msa_repr = self.msa_transformer(msa_emb)  # [batch_size, num_seqs, seq_len, d_model]\n        \n        # 쿼리 시퀀스 (첫 번째 시퀀스)에 대한 임베딩\n        query_repr = msa_repr[:, 0]  # [batch_size, seq_len, d_model]\n        query_repr = self.single_seq_embedding(query_repr)\n        \n        # MSA 정보를 쿼리 시퀀스에 크로스 어텐션으로 통합\n        query_refined = self.cross_transformer(query_repr, msa_repr)\n        \n        # 구조 모듈로 좌표 예측\n        xyz_pred = self.structure_module(query_refined.unsqueeze(1)).squeeze(1)\n        \n        return xyz_pred\n\n# 추론 함수 - 앙상블 예측 생성\ndef predict_ensemble(model, batch, num_ensemble=5):\n    \"\"\"\n    드롭아웃을 활성화한 상태에서 여러 번 예측하여 앙상블 결과 생성\n    \"\"\"\n    model.train()  # 드롭아웃 활성화\n    \n    sequence = batch['sequence'].cuda()\n    msa = batch['msa'].cuda()\n    \n    predictions = []\n    \n    for _ in range(num_ensemble):\n        with torch.no_grad():\n            pred_xyz = model(sequence, msa).squeeze(0).cpu().numpy()\n            predictions.append(pred_xyz)\n    \n    return predictions\n","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-03-15T07:59:44.251556Z","iopub.execute_input":"2025-03-15T07:59:44.251997Z","iopub.status.idle":"2025-03-15T07:59:44.294210Z","shell.execute_reply.started":"2025-03-15T07:59:44.251973Z","shell.execute_reply":"2025-03-15T07:59:44.293432Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"set_seed(42)\n\n# 테스트 데이터 로드\ntest_data = pd.read_csv(\"/kaggle/input/stanford-rna-3d-folding/test_sequences.csv\")\nprint(f\"테스트 데이터 로드 완료: {len(test_data)} 샘플\")\n\n# MSA 디렉토리\nmsa_dir = \"/kaggle/input/stanford-rna-3d-folding/MSA\"\n\n# 데이터셋 생성\ntest_dataset = RNA3D_MSA_TestDataset(test_data, msa_dir, config)\ntest_loader = DataLoader(test_dataset, batch_size=1, shuffle=False)\n\n# 모델 로드\nmodel = RNA_MSA_Folding(config).cuda()\nmodel.load_state_dict(torch.load('/kaggle/input/20250315-rna-msa-folding/RNA_MSA_Folding_best.pt'))\nprint(\"모델 로드 완료\")\n\n# 예측 실행\nall_predictions = []\ntarget_ids = []\nraw_sequences = []\n\nfor batch in tqdm(test_loader, desc=\"예측 중\"):\n    target_id = batch['target_id'][0]\n    raw_sequence = batch['raw_sequence'][0]\n    \n    # 앙상블 예측\n    predictions = predict_ensemble(model, batch, num_ensemble=config['num_ensemble'])\n    \n    all_predictions.append(predictions)\n    target_ids.append(target_id)\n    raw_sequences.append(raw_sequence)\n\n# 제출 파일 형식으로 변환\ndata = []\n\nfor i in range(len(target_ids)):\n    for j in range(len(raw_sequences[i])):\n        # ID, resname, resid\n        row = [f\"{target_ids[i]}_{j+1}\", raw_sequences[i][j], j+1]\n        \n        # 5개 앙상블 모델의 x, y, z 좌표 추가\n        for k in range(config['num_ensemble']):\n            for coord_idx in range(3):  # x, y, z\n                row.append(all_predictions[i][k][j][coord_idx])\n        \n        data.append(row)\n\n# 열 이름 생성\ncolumns = ['ID', 'resname', 'resid']\nfor i in range(1, config['num_ensemble'] + 1):\n    columns.extend([f\"x_{i}\", f\"y_{i}\", f\"z_{i}\"])\n\n# 데이터프레임 생성 및 저장\nsubmission = pd.DataFrame(data, columns=columns)\nsubmission.to_csv('submission.csv', index=False)\n\nprint(f\"제출 파일 생성 완료: {len(submission)} 행\")\nprint(submission.head())","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-03-15T07:59:44.295136Z","iopub.execute_input":"2025-03-15T07:59:44.295411Z","iopub.status.idle":"2025-03-15T07:59:47.444221Z","shell.execute_reply.started":"2025-03-15T07:59:44.295390Z","shell.execute_reply":"2025-03-15T07:59:47.443303Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"","metadata":{"trusted":true},"outputs":[],"execution_count":null}]}