File size: 1,508 Bytes
52f4aa2 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 | #!/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]}")
|