from transformers import AutoModelForCausalLM, AutoTokenizer, TextGenerationPipeline

class HuggingFaceHandler:
    def __init__(self, model_dir):
        """
        Initialize the handler with the model directory.
        """
        # Charger le tokenizer
        self.tokenizer = AutoTokenizer.from_pretrained(model_dir)
        
        # Charger le modèle de génération de texte
        self.model = AutoModelForCausalLM.from_pretrained(model_dir)
        
        # Initialiser le pipeline de génération de texte
        self.pipeline = TextGenerationPipeline(model=self.model, tokenizer=self.tokenizer, framework='pt')

    def __call__(self, subject):
        """
        Generate a math course based on the given subject.
        """
        # Générer le texte à partir du sujet
        generated_text = self.pipeline(subject, max_length=500)  # Ajuster max_length selon les besoins
        
        return generated_text

# Assumer que le chemin du modèle est déjà spécifié lors de l'initialisation de l'objet handler.
model_path = "KakaoL0L/Mistral7B_MatheoAI_lora"

# Créer une instance du gestionnaire
handler = HuggingFaceHandler(model_path)