File size: 1,170 Bytes
81f487c 5969acd 81f487c |
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 |
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) |