

import pandas as pd
train = pd.read_csv('train_data.csv')
test= pd.read_csv("./test_data.csv")
train.head()
temp = test['sequence'].copy()
train['target_profile'] = train['target_profile'].str.split(" ")
len(train['target_profile'][1])
train['sequence'] = train['sequence'].apply(list)
test['sequence'] = test['sequence'].apply(list)
train['sequence']

import numpy as np
t = np.zeros(256-len(train['sequence'][1]))
print(t)
for i in range(len(train)):
    seq = np.array(train['sequence'][i])
    pad = np.ones(256 - len(seq)) * -1.0
    train.at[i,'sequence'] = np.hstack([seq,pad])
for i in range(len(test)):
    seq = np.array(test['sequence'][i])
    pad = np.ones(256 - len(seq)) * -1.0
    test.at[i,'sequence'] = np.hstack([seq,pad])
mapping_seq = { 'A': 1,'C': 2,'G': 3,'U': 4,'-1.0': 0}
train['sequence'] = train['sequence'].apply(lambda seq: [mapping_seq[x] for x in seq])
test['sequence'] = test['sequence'].apply(lambda seq: [mapping_seq[x] for x in seq])
train['sequence'][10]
for i in range(len(train)):
    seq = np.array(train['target_profile'][i])
    pad = np.ones(256 - len(seq)) * -1.0
    train.at[i,'target_profile'] = np.hstack([seq,pad])
train['target_profile'] = train['target_profile'].apply(lambda x: [float(v) for v in x])
train['target_profile'][10]
import torch
import torch.nn as nn
import torch.optim as optim
from torchvision import models, transforms
from torch.utils.data import DataLoader,Dataset,random_split
class TrainDataset(Dataset):
    def __init__(self, train):
        self.df = train.reset_index(drop=True)
    def __len__(self):
        return len(self.df)
    def __getitem__(self, idx):
        seq = torch.tensor(self.df.loc[idx, 'sequence'], dtype=torch.long)
        target = torch.tensor(self.df.loc[idx, 'target_profile'], dtype=torch.float32)
        mask = (target != -1.0).float()
        return seq, target, mask
class LSTM(nn.Module):
    def __init__(self):
        super().__init__()
        self.emb = nn.Embedding(5, 16, padding_idx=0)
        self.lstm = nn.LSTM(16,64,num_layers=3,batch_first=True,bidirectional=True)
        self.fc = nn.Linear(64 * 2, 1)
    def forward(self, x):
        x = self.emb(x)
        x, _ = self.lstm(x)
        x = self.fc(x).squeeze(-1)
        return x
train_ds = TrainDataset(train)
train_loader = DataLoader(train_ds, batch_size=32, shuffle=True, num_workers=0)
device = torch.device('cuda')
model = LSTM()
model = model.to(device)
optimizer = optim.Adam(model.parameters(), lr=1e-4)
criterion = nn.MSELoss(reduction='none')
scheduler = optim.lr_scheduler.StepLR(optimizer, step_size=30, gamma=0.2)
for epoch in range(20):
    total_loss = 0
    for seq, target, mask in train_loader:
        seq = seq.to(device)
        target= target.to(device)
        mask = mask.to(device)
        optimizer.zero_grad()
        output = model(seq)
        loss = (criterion(output, target) * mask).sum() / mask.sum()
        loss.backward()
        optimizer.step()
        total_loss += loss.item()
    scheduler.step()
    print(f"epoch {epoch+1}, loss {total_loss/len(train_loader)}")
class TestDataset(Dataset):
    def __init__(self, test):
        self.df = test.reset_index(drop=True)
    def __len__(self):
        return len(self.df)
    def __getitem__(self, idx):
        seq = torch.tensor(self.df.loc[idx, 'sequence'], dtype=torch.long)
        return seq
test['sequence'][1]
test_ds = TestDataset(test)
test_loader= DataLoader(test_ds, shuffle=False, num_workers=0, batch_size=32)
preds = []
for seq in test_loader:
    with torch.no_grad():
        seq = seq.to(device)
        output = model(seq)
        preds.append(output.cpu())
preds_final = torch.cat(preds)
preds_final = preds_final.numpy()
preds_final
ttt = ""
for t in preds_final[1]:
    ttt+=str(t)
    ttt+=" "
ttt
preds_reale_final = []
for i in range(len(preds_final)):
    L = len(temp.iloc[i])
    s = " ".join(str(x) for x in preds_final[i][:L])
    preds_reale_final.append(s)
subs = pd.DataFrame({
    'subtaskID':1,
    'datapointID':test['id'],
    'answer':preds_reale_final
}).to_csv('subs.csv',index=False)
