| import json | |
| from transformers import AutoTokenizer, AutoModelForCausalLM | |
| # Nom du modèle que vous souhaitez utiliser | |
| model_name = "AIDC-AI/Ovis1.6-Gemma2-9B" # Remplacez par le nom de votre modèle | |
| # Initialisation globale des variables | |
| tokenizer = None | |
| model = None | |
| # Fonction d'initialisation qui sera appelée lors du démarrage du service | |
| def init(): | |
| global tokenizer, model | |
| tokenizer = AutoTokenizer.from_pretrained(model_name, trust_remote_code=True) | |
| model = AutoModelForCausalLM.from_pretrained(model_name, trust_remote_code=True) | |
| model.eval() # Définir le modèle en mode évaluation | |
| # Fonction de prédiction qui sera appelée pour traiter les requêtes d'inférence | |
| def predict(data): | |
| global tokenizer, model | |
| # Extraire les données d'entrée | |
| inputs = data.get("inputs") | |
| # Vérifier si les données d'entrée sont une chaîne de caractères | |
| if isinstance(inputs, str): | |
| inputs = tokenizer(inputs, return_tensors="pt") | |
| # Générer les prédictions à partir du modèle | |
| outputs = model.generate(**inputs) | |
| # Convertir les résultats en texte | |
| result = tokenizer.decode(outputs[0], skip_special_tokens=True) | |
| # Retourner le résultat au format JSON | |
| return json.dumps({"result": result}) | |
| # Fonction d'initialisation qui sera appelée lors du démarrage du service | |
| init() | |
| # Fonction de prédiction qui sera appelée pour traiter les requêtes d'inférence | |
| def main(): | |
| data = {"inputs": "Votre texte d'entrée"} # Remplacez par les données d'entrée | |
| result = predict(data) | |
| print(result) | |
| if __name__ == "__main__": | |
| main() |