#!/usr/bin/env python3
"""
v4 Fold0 推論スクリプト
"""

import os
import pandas as pd
import torch
from transformers import AutoTokenizer, AutoModelForSeq2SeqLM
from tqdm import tqdm

# パス設定
MODEL_PATH = "/kaggle/input/byt5-v4-fold0"
TEST_PATH = "/kaggle/input/deep-past-initiative-machine-translation/test.csv"
OUTPUT_PATH = "/kaggle/working/submission.csv"

# 設定
PREFIX = "translate Akkadian to English: "
MAX_LENGTH = 512
BATCH_SIZE = 8

# 生成パラメータ（v4と同じ）
GEN_PARAMS = {
    "max_length": 512,
    "num_beams": 4,
    "repetition_penalty": 1.2,
    "length_penalty": 0.8,
    "early_stopping": True
}

print(f"Model path: {MODEL_PATH}")
print(f"Files in model path: {os.listdir(MODEL_PATH)}")

# モデルとトークナイザーのロード
print("Loading model...")
tokenizer = AutoTokenizer.from_pretrained(MODEL_PATH)
model = AutoModelForSeq2SeqLM.from_pretrained(MODEL_PATH)

device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
print(f"Device: {device}")
model.to(device)
model.eval()

# テストデータ読み込み
print("Loading test data...")
test_df = pd.read_csv(TEST_PATH)
print(f"Test samples: {len(test_df)}")

# 推論
results = []
for i in tqdm(range(0, len(test_df), BATCH_SIZE)):
    batch = test_df.iloc[i:i+BATCH_SIZE]
    texts = [PREFIX + str(row['transliteration']) for _, row in batch.iterrows()]
    
    inputs = tokenizer(
        texts,
        return_tensors="pt",
        padding=True,
        truncation=True,
        max_length=MAX_LENGTH
    ).to(device)
    
    with torch.no_grad():
        outputs = model.generate(**inputs, **GEN_PARAMS)
    
    translations = tokenizer.batch_decode(outputs, skip_special_tokens=True)
    
    for idx, (_, row) in enumerate(batch.iterrows()):
        results.append({
            "id": row['id'],
            "translation": translations[idx]
        })

# 提出ファイル作成
submission_df = pd.DataFrame(results)
submission_df.to_csv(OUTPUT_PATH, index=False)
print(f"Submission saved: {len(submission_df)} rows")
print(submission_df.head())