import torch
import sys
import os

MODEL_PATH = "/workspace/ankahi/punjabi_gemma/ankahi"

print(f"Python version: {sys.version}")
print(f"PyTorch version: {torch.__version__}")
print(f"CUDA available: {torch.cuda.is_available()}")

if torch.cuda.is_available():
    print(f"CUDA version: {torch.version.cuda}")
    print(f"Device count: {torch.cuda.device_count()}")
    for i in range(torch.cuda.device_count()):
        print(f"Device {i}: {torch.cuda.get_device_name(i)}")

try:
    from transformers import AutoModelForCausalLM, AutoProcessor
    print("Transformers imported.")
    
    print("Loading processor...")
    processor = AutoProcessor.from_pretrained(MODEL_PATH)
    
    print("Loading model on CPU first to check if weights are readable...")
    # This might take a lot of RAM
    model = AutoModelForCausalLM.from_pretrained(
        MODEL_PATH,
        torch_dtype=torch.bfloat16,
        low_cpu_mem_usage=True,
        device_map="auto"
    )
    print("Model loaded successfully!")
except Exception as e:
    print(f"Caught exception: {e}")
    import traceback
    traceback.print_exc()