#!/usr/bin/env python3
"""
Example inference script for CPU deployment
"""
import torch
from transformers import AutoTokenizer
from pathlib import Path

# Charger le modèle
def load_model(model_path, quantized=True):
    """Charge le modèle pour l'inférence"""
    if quantized:
        checkpoint = torch.load(model_path / 'model_quantized.pt', map_location='cpu')
    else:
        checkpoint = torch.load(model_path / 'model_fp32.pt', map_location='cpu')

    # Recréer le modèle
    from src.models.student_model import LEAFStudent
    from transformers import AutoConfig

    # Charger config depuis le checkpoint
    model = LEAFStudent(
        teacher_config=checkpoint['config'],
        pooling_mode=checkpoint.get('pooling_mode', 'mean')
    )

    model.load_state_dict(checkpoint['model_state_dict'], strict=False)
    model.eval()

    # Charger tokenizer
    tokenizer = AutoTokenizer.from_pretrained(model_path)
    model.set_tokenizer(tokenizer)

    return model

# Utilisation
if __name__ == "__main__":
    model_dir = Path(__file__).parent

    # Charger modèle quantized (plus rapide)
    model = load_model(model_dir, quantized=True)

    # Test
    texts = [
        "This is a test sentence",
        "Machine learning is awesome",
    ]

    with torch.no_grad():
        embeddings = model.encode(texts, device='cpu')

    print(f"Embeddings shape: {embeddings.shape}")
    print(f"Sample embedding: {embeddings[0][:5]}")