# Akkadian Translation - T5 Inference (v2)
# Model weights from akkadian-t5-model, tokenizer from t5-base-tokenizer
# Improved generation settings

import pandas as pd
import numpy as np
from transformers import T5ForConditionalGeneration, T5Tokenizer
import torch
import os

# Configuration - separate paths for model and tokenizer
MODEL_PATH = "/kaggle/input/akkadian-t5-model"
TOKENIZER_PATH = "/kaggle/input/t5-base-tokenizer"
TEST_PATH = "/kaggle/input/deep-past-initiative-machine-translation/test.csv"
OUTPUT_PATH = "/kaggle/working/submission.csv"

print("Model files:")
for f in os.listdir(MODEL_PATH):
    size = os.path.getsize(f"{MODEL_PATH}/{f}")
    print(f"  {f}: {size/1024:.1f}KB")

print("\nTokenizer files:")
for f in os.listdir(TOKENIZER_PATH):
    size = os.path.getsize(f"{TOKENIZER_PATH}/{f}")
    print(f"  {f}: {size/1024:.1f}KB")

print("\nLoading tokenizer from separate dataset...")
tokenizer = T5Tokenizer.from_pretrained(TOKENIZER_PATH, local_files_only=True)

print("Loading model weights...")
model = T5ForConditionalGeneration.from_pretrained(MODEL_PATH, local_files_only=True)

# Move to GPU
device = "cuda" if torch.cuda.is_available() else "cpu"
model = model.to(device)
model.eval()
print(f"Model loaded on {device}")

# Load test data
print("\nLoading test data...")
test_df = pd.read_csv(TEST_PATH)
print(f"Test samples: {len(test_df)}")
print(test_df)

# Generate predictions with improved settings
print("\nGenerating predictions with improved beam search...")
predictions = []

with torch.no_grad():
    for idx, row in test_df.iterrows():
        input_text = f"translate Akkadian to English: {row['transliteration']}"
        inputs = tokenizer(
            input_text,
            return_tensors="pt",
            max_length=512,
            truncation=True,
            padding=True
        ).to(device)
        
        # Improved generation settings
        outputs = model.generate(
            **inputs,
            max_length=512,  # Longer output
            min_length=10,   # Minimum length
            num_beams=8,     # More beams
            early_stopping=True,
            no_repeat_ngram_size=3,
            length_penalty=1.0,
            repetition_penalty=1.2,  # Discourage repetition
        )
        
        translation = tokenizer.decode(outputs[0], skip_special_tokens=True)
        predictions.append(translation)
        print(f"ID {row['id']}: {translation[:100]}...")

# Create submission
submission_df = pd.DataFrame({
    "id": test_df["id"],
    "translation": predictions
})

submission_df.to_csv(OUTPUT_PATH, index=False)
print(f"\nSubmission saved to {OUTPUT_PATH}")
print("\n=== Final Submission ===")
print(submission_df)