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

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

# 接頭辞（必須）
PREFIX = "translate Akkadian to English: "

print(f"Loading model from {MODEL_PATH}")
print(f"Files in MODEL_PATH: {os.listdir(MODEL_PATH)}")

# モデルとトークナイザーのロード（ローカルファイルから）
tokenizer = AutoTokenizer.from_pretrained(MODEL_PATH, local_files_only=True)
model = AutoModelForSeq2SeqLM.from_pretrained(MODEL_PATH, local_files_only=True)
model.eval()

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

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

results = []

# 推論
batch_size = 8
for i in tqdm(range(0, len(test_df), batch_size), desc="Inference"):
    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=256
    ).to(device)
    
    with torch.no_grad():
        outputs = model.generate(
            **inputs,
            max_length=256,
            num_beams=5,
            no_repeat_ngram_size=3,
            early_stopping=True
        )
    
    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())