agent-ai / core /model_loader.py
Mauricio-100's picture
Create core/model_loader.py
8775924 verified
raw
history blame contribute delete
278 Bytes
from transformers import AutoModelForCausalLM
import torch
def load_model(model_name: str):
print(f"🧠 Chargement du modèle {model_name}...")
model = AutoModelForCausalLM.from_pretrained(
model_name,
torch_dtype=torch.bfloat16
)
return model