# %%
# This Python 3 environment comes with many helpful analytics libraries installed
# It is defined by the kaggle/python Docker image: https://github.com/kaggle/docker-python
# For example, here's several helpful packages to load

import numpy as np # linear algebra
import pandas as pd # data processing, CSV file I/O (e.g. pd.read_csv)

# Input data files are available in the read-only "../input/" directory
# For example, running this (by clicking run or pressing Shift+Enter) will list all files under the input directory

import os
for dirname, _, filenames in os.walk('/kaggle/input'):
    for filename in filenames:
        print(os.path.join(dirname, filename))

# You can write up to 20GB to the current directory (/kaggle/working/) that gets preserved as output when you create a version using "Save & Run All" 
# You can also write temporary files to /kaggle/temp/, but they won't be saved outside of the current session

# Use the kagglehub client library to attach Kaggle resources like competitions, datasets, and models to your session
# Learn more about kagglehub: https://github.com/Kaggle/kagglehub/blob/main/README.md

import kagglehub
# kagglehub.dataset_download('<owner>/<dataset-slug>')
# %%
import pandas as pd
import torch
import torch.nn as nn
from transformers import AutoTokenizer, AutoModelForSequenceClassification
from torch.utils.data import Dataset, DataLoader
# %%
train_df = pd.read_csv('/kaggle/input/competitions/nlp-getting-started/train.csv')
test_df = pd.read_csv('/kaggle/input/competitions/nlp-getting-started/test.csv')
from sklearn.model_selection import train_test_split
train_df, val_df = train_test_split(train_df, test_size=0.1, random_state=42)
print(f"Train: {len(train_df)}, Val: {len(val_df)}, Test: {len(test_df)}")
# %%
model_name = "xlm-roberta-large"
tokenizer = AutoTokenizer.from_pretrained(model_name)
model = AutoModelForSequenceClassification.from_pretrained(model_name, num_labels=2)
# %%
import numpy as np
class NliDataset(Dataset):
    def __init__(self, df, tokenizer, max_length=128):
        self.df = df.reset_index(drop=True)
        self.tokenizer = tokenizer
        self.max_length = max_length

    def __len__(self):
        return len(self.df)

    def __getitem__(self, idx):
        row = self.df.iloc[idx]
        strs=''
        if row['keyword'] is not np.nan:
            strs=strs+'['+row['keyword']+']'
        if row['location'] is not np.nan:
            strs=strs+'['+row['location']+']'
        strs=strs+row['text']
        encoding = self.tokenizer(
            strs,
            truncation=True,
            padding='max_length',
            max_length=self.max_length,
            return_tensors='pt'  # 返回 PyTorch Tensor
        )
        # 去掉 batch 维度（因为只处理单条）
        item = {
            'input_ids': encoding['input_ids'].squeeze(0),
            'attention_mask': encoding['attention_mask'].squeeze(0)
        }
        if 'target' in row:
            item['labels'] = torch.tensor(row['target'], dtype=torch.long)
        return item

batch_size = 16  # 根据显存调整，推荐使用 16 或 32
max_length = 128

train_dataset = NliDataset(train_df, tokenizer, max_length)
val_dataset = NliDataset(val_df, tokenizer, max_length)

train_loader = DataLoader(train_dataset, batch_size=batch_size, shuffle=True)
val_loader = DataLoader(val_dataset, batch_size=batch_size, shuffle=False)
# %%
for name, param in model.named_parameters():
    #if param.requires_grad:
    print(f"{name}: requires_grad = {param.requires_grad}")
# %%
from peft import LoraConfig, get_peft_model
target_layers = [8,9,10,11,12,13,14,15,16,17,18,19,20,21,22,23]
target_modules=[]
for layer in target_layers:
    target_modules.append(f"roberta.encoder.layer.{layer}.attention.self.query")
    target_modules.append(f"roberta.encoder.layer.{layer}.attention.self.value")
# 配置 LoRA
config = LoraConfig(
    r=8,  # LoRA 的秩
    lora_alpha=16,  # LoRA 的缩放因子
    target_modules=target_modules,  # 目标模块
    lora_dropout=0.1,  # Dropout 概率
    bias="none",  # 是否更新偏置
    modules_to_save=["classifier"],  # 指定分类器需要被微调
)

# 封装为 LoRA 模型
model = get_peft_model(model, config)

# 验证分类器是否被微调
print("验证分类器参数是否被训练:")
for name, param in model.named_parameters():
    if param.requires_grad:
        print(f"{name}: requires_grad = {param.requires_grad}")
# %%
@torch.no_grad()
def validate(model,loader):
    model.eval()
    acc=0
    total=0
    for batch in loader:
        input_ids = batch['input_ids'].to(device)
        attention_mask = batch['attention_mask'].to(device)
        labels = batch['labels'].to(device)
        outputs = model(input_ids, attention_mask=attention_mask).logits
        pred=torch.argmax(outputs,dim=1)
        acc+=pred.eq(labels).sum()
        total+=labels.size(0)
    print(f"acc:{acc/total}")
    return acc/total
from tqdm import tqdm   
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
model=model.to(device)
loss_func = nn.CrossEntropyLoss()
optimizer = torch.optim.AdamW(model.parameters(), lr=2e-4)
scheduler=torch.optim.lr_scheduler.CosineAnnealingLR(optimizer, T_max=10)
epochs = 30
best_acc=0.0
for epoch in range(epochs):
    model.train()
    training_loss = 0
    
    # 使用 tqdm 包装 dataloader，并设置描述信息
    progress_bar = tqdm(train_loader, desc=f"Epoch {epoch+1}/{epochs}")
    lens=0
    for batch in progress_bar:
        optimizer.zero_grad()
        input_ids = batch['input_ids'].to(device)
        attention_mask = batch['attention_mask'].to(device)
        labels = batch['labels'].to(device)
        outputs = model(input_ids, attention_mask=attention_mask).logits
        loss = loss_func(outputs, labels)
        loss.backward()
        optimizer.step()
        
        training_loss += loss.item()
        lens+=1
        # 更新进度条显示当前 batch 的损失
        progress_bar.set_postfix({
            'loss': f'{loss.item():.4f}',
            'avg_loss': f'{training_loss / (progress_bar.n+1):.4f}'  # progress_bar.n 是已处理 batch 数
        })
    
    scheduler.step()
    
    avg_train_loss = training_loss / lens
    print(f"Epoch {epoch+1} train_loss: {avg_train_loss:.4f}")
    
    # 验证（你也可以为验证添加进度条，见下方建议）
    current_acc=validate(model, val_loader)
    if current_acc > best_acc :
        torch.save(model.state_dict(), 'model.pth')
        print(f"best model save,acc:{current_acc}")
        best_acc=current_acc
# %%
test_dataset = NliDataset(test_df, tokenizer, max_length)
test_loader = DataLoader(test_dataset, batch_size=batch_size, shuffle=False)
all_preds=[]
model.load_state_dict(torch.load('model.pth'))
model.eval()
with torch.no_grad():
    for batch in test_loader:
        input_ids = batch['input_ids'].to(device)
        attention_mask = batch['attention_mask'].to(device)
        outputs = model(input_ids, attention_mask=attention_mask).logits
        preds = torch.argmax(outputs, dim=1)
        all_preds.extend(preds.cpu().numpy())

submission = pd.DataFrame({
    'id':test_df['id'],
    'target': all_preds
})

print(submission)
submission.to_csv('submission.csv', index=False)
print("Submission saved!")
# %%