Ajout du fichier handler.py

handler.py

@@ -0,0 +1,26 @@
+import json
+import torch
+from transformers import AutoModelForCausalLM, AutoTokenizer
+
+# Charger le modèle et le tokenizer
+model_name = "AIDC-AI/Ovis1.6-Gemma2-9B"  # Remplacez par le nom de votre modèle
+tokenizer = AutoTokenizer.from_pretrained(model_name, trust_remote_code=True)
+model = AutoModelForCausalLM.from_pretrained(model_name, trust_remote_code=True)
+
+# Fonction d'initialisation
+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
+def predict(data):
+	inputs = data.get("inputs")
+	if isinstance(inputs, str):
+		inputs = tokenizer(inputs, return_tensors="pt")
+
+	outputs = model.generate(**inputs)
+	result = tokenizer.decode(outputs[0], skip_special_tokens=True)
+	
+	return json.dumps({"result": result})