# Akkadian Translation - M2M-100 Inference
# Model trained on Modal, achieving better multilingual translation

import pandas as pd
import numpy as np
from transformers import M2M100ForConditionalGeneration, M2M100Tokenizer
import torch
import os

# Configuration
MODEL_PATH = "/kaggle/input/akkadian-m2m-model"
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/1024:.1f}MB" if size > 1024*1024 else f"  {f}: {size/1024:.1f}KB")

print("\nLoading tokenizer...")
tokenizer = M2M100Tokenizer.from_pretrained(MODEL_PATH, local_files_only=True)

print("Loading model weights...")
model = M2M100ForConditionalGeneration.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)

# Set language codes for inference
# Using Indonesian (id) as proxy for Akkadian during training
tokenizer.src_lang = "id"
forced_bos_token_id = tokenizer.lang_code_to_id["en"]

# Generate predictions
print("\nGenerating predictions...")
predictions = []

with torch.no_grad():
    for idx, row in test_df.iterrows():
        input_text = row['transliteration']
        
        inputs = tokenizer(
            input_text,
            return_tensors="pt",
            max_length=256,
            truncation=True,
            padding=True
        ).to(device)
        
        outputs = model.generate(
            inputs.input_ids,
            attention_mask=inputs.attention_mask,
            forced_bos_token_id=forced_bos_token_id,
            max_length=512,
            num_beams=5,
            early_stopping=True,
            no_repeat_ngram_size=3,
            length_penalty=1.0,
            do_sample=False,
        )
        
        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)