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

# モデルパス (v6)
MODEL_PATH = "/kaggle/input/byt5-v6-checkpoint"

print("Loading model (v6 checkpoint-10000)...")
print(f"Files in MODEL_PATH: {os.listdir(MODEL_PATH)}")

tokenizer = AutoTokenizer.from_pretrained(MODEL_PATH)
model = AutoModelForSeq2SeqLM.from_pretrained(MODEL_PATH)
model.eval()

if torch.cuda.is_available():
    model = model.cuda()
    print("Using GPU")
else:
    print("Using CPU")

# テストデータ読み込み
test = pd.read_csv("/kaggle/input/deep-past-initiative-machine-translation/test.csv")
print(f"Test data: {len(test)} rows")

# プレフィックス
PREFIX = "translate Akkadian to English: "

# 推論
predictions = []
batch_size = 8

for i in range(0, len(test), batch_size):
    batch = test.iloc[i:i+batch_size]
    sources = [PREFIX + str(row["transliteration"]) for _, row in batch.iterrows()]
    
    inputs = tokenizer(
        sources,
        max_length=512,
        truncation=True,
        padding=True,
        return_tensors="pt"
    )
    
    if torch.cuda.is_available():
        inputs = {k: v.cuda() for k, v in inputs.items()}
    
    with torch.no_grad():
        outputs = model.generate(
            **inputs,
            max_length=512,
            num_beams=4,
            early_stopping=True,
            repetition_penalty=1.2,
            length_penalty=0.8
        )
    
    decoded = tokenizer.batch_decode(outputs, skip_special_tokens=True)
    predictions.extend(decoded)
    
    if (i // batch_size) % 50 == 0:
        print(f"Progress: {i}/{len(test)}")

print(f"Total predictions: {len(predictions)}")

# 提出ファイル作成
submission = pd.DataFrame({
    "id": test["id"],
    "translation": predictions
})

submission.to_csv("submission.csv", index=False)
print("Submission saved!")
print(submission.head())