|
|
| """
|
| Created on Sun Sep 15 18:27:17 2024
|
|
|
| @author: salikha4
|
| """
|
|
|
| import os
|
| import csv
|
| import json
|
| import shutil
|
| import random
|
| import argparse
|
| from datetime import datetime
|
| import pandas as pd
|
| import time
|
| import torch
|
| import torch.nn as nn
|
| import torch.nn.functional as F
|
| from torch.utils.data import Dataset, DataLoader, TensorDataset
|
| from torch.optim import Adam
|
| import numpy as np
|
| import warnings
|
| warnings.filterwarnings('ignore')
|
|
|
| def lwm_inference(preprocessed_chs, input_type, lwm_model, device):
|
|
|
| dataset = prepare_for_lwm(preprocessed_chs, device)
|
|
|
| lwm_loss, embedding_data = evaluate(lwm_model, dataset)
|
|
|
|
|
| if input_type == 'cls_emb':
|
| embedding_data = embedding_data[:, 0]
|
| elif input_type == 'channel_emb':
|
| embedding_data = embedding_data[:, 1:]
|
|
|
| dataset = embedding_data.float()
|
| return dataset
|
|
|
| def prepare_for_lwm(data, device, batch_size=64, shuffle=False):
|
|
|
| input_ids, masked_tokens, masked_pos = zip(*data)
|
|
|
| input_ids_tensor = torch.tensor(input_ids, device=device).float()
|
| masked_tokens_tensor = torch.tensor(masked_tokens, device=device).float()
|
| masked_pos_tensor = torch.tensor(masked_pos, device=device).long()
|
|
|
| dataset = TensorDataset(input_ids_tensor, masked_tokens_tensor, masked_pos_tensor)
|
|
|
| return DataLoader(dataset, batch_size=batch_size, shuffle=shuffle)
|
|
|
| def evaluate(model, dataloader):
|
|
|
| model.eval()
|
| running_loss = 0.0
|
| outputs = []
|
| criterionMCM = nn.MSELoss()
|
|
|
| with torch.no_grad():
|
| for idx, batch in enumerate(dataloader):
|
| input_ids = batch[0]
|
| masked_tokens = batch[1]
|
| masked_pos = batch[2]
|
|
|
| logits_lm, output = model(input_ids, masked_pos)
|
|
|
| output_batch_preproc = output
|
| outputs.append(output_batch_preproc)
|
|
|
| loss_lm = criterionMCM(logits_lm, masked_tokens)
|
| loss = loss_lm / torch.var(masked_tokens)
|
| running_loss += loss.item()
|
|
|
| average_loss = running_loss / len(dataloader)
|
| output_total = torch.cat(outputs, dim=0)
|
|
|
| return average_loss, output_total
|
|
|
| def create_raw_dataset(data, device):
|
| """Create a dataset for raw channel data."""
|
| input_ids, _, _ = zip(*data)
|
| input_data = torch.tensor(input_ids, device=device)[:, 1:]
|
| return input_data.float()
|
| |
