def train(args):
    tokenizer = AutoTokenizer.from_pretrained(Params.pretrained_model_path)
    dl = read_data(args)
    model = SimCSE(Params.pretrained_model_path, args.pool_type, args.dropout_rate).to(args.device)
    optimizer = torch.optim.AdamW(model.parameters(), lr=args.lr)
    model.train()
    batch_idx = 0
    min_loss = 10000000
    for epoch_idx in range(args.epochs):
        epoch_losses = []
        for data in tqdm(dl):
            batch_idx += 1
            new_batch_data = duplicate_batch(data, tokenizer, args)
            pred = model(input_ids=new_batch_data['input_ids'].to(args.device), attention_mask=new_batch_data['attention_mask'].to(args.device), token_type_ids=new_batch_data['token_type_ids'].to(args.device))
            loss = compute_loss(pred, args.tao, args.device)
            optimizer.zero_grad()
            loss.backward()
            optimizer.step()
            loss = loss.item()
            epoch_losses.append(loss)
            if ((batch_idx % args.display_interval) == 0):
                logging.info(f'epoch: {epoch_idx}, batch_idx: {batch_idx}, loss: {loss:>10f}')
        avg_epoch_loss = np.mean(epoch_losses)
        if (avg_epoch_loss < min_loss):
            min_loss = avg_epoch_loss
            torch.save({'epoch': epoch_idx, 'model_state_dict': model.state_dict(), 'loss': avg_epoch_loss}, Params.simcse_model_path)
def train(args):
    tokenizer = AutoTokenizer.from_pretrained(Params.pretrained_model_path)
    dl = read_data(args)
    model = SimCSE(Params.pretrained_model_path, args.pool_type, args.dropout_rate).to(args.device)
    optimizer = torch.optim.AdamW(model.parameters(), lr=args.lr)
    model.train()
    batch_idx = 0
    min_loss = 10000000
    for epoch_idx in range(args.epochs):
        epoch_losses = []
        for data in tqdm(dl):
            batch_idx += 1
            new_batch_data = duplicate_batch(data, tokenizer, args)
            pred = model(input_ids=new_batch_data['input_ids'].to(args.device), attention_mask=new_batch_data['attention_mask'].to(args.device), token_type_ids=new_batch_data['token_type_ids'].to(args.device))
            loss = compute_loss(pred, args.tao, args.device)
            optimizer.zero_grad()
            loss.backward()
            optimizer.step()
            loss = loss.item()
            epoch_losses.append(loss)
            if ((batch_idx % args.display_interval) == 0):
                logging.info(f'epoch: {epoch_idx}, batch_idx: {batch_idx}, loss: {loss:>10f}')
        avg_epoch_loss = np.mean(epoch_losses)
        if (avg_epoch_loss < min_loss):
            min_loss = avg_epoch_loss
            torch.save({'epoch': epoch_idx, 'model_state_dict': model.state_dict(), 'loss': avg_epoch_loss}, Params.simcse_model_path)