Update app.py

app.py

@@ -1,10 +1,11 @@
 # Importer les bibliothèques nécessaires
 import gradio as gr
-from transformers import NVLM_D, AutoTokenizer
+from transformers import AutoModelForCausalLM, AutoTokenizer
 
 # Charger le modèle et le tokenizer
-model = NVLM_D.from_pretrained("SeanScripts/NVLM-D-72B-nf4", trust_remote_code=True)
-tokenizer = AutoTokenizer.from_pretrained("SeanScripts/NVLM-D-72B-nf4", trust_remote_code=True)
+model_name = "SeanScripts/NVLM-D-72B-nf4"
+model = AutoModelForCausalLM.from_pretrained(model_name, trust_remote_code=True)
+tokenizer = AutoTokenizer.from_pretrained(model_name, trust_remote_code=True)
 
 # Fonction pour effectuer des prédictions
 def predict(text):
@@ -12,10 +13,10 @@ def predict(text):
     inputs = tokenizer(text, return_tensors="pt")
     
     # Effectuer la prédiction
-    outputs = model(**inputs)
+    outputs = model.generate(**inputs)
     
-    # Extraire et retourner les résultats (vous pouvez adapter cela selon vos besoins)
-    return outputs.logits.argmax(dim=-1).item()  # Exemple de retour de la classe prédite
+    # Décoder les résultats et retourner le texte généré
+    return tokenizer.decode(outputs[0], skip_special_tokens=True)
 
 # Créer l'interface Gradio
 iface = gr.Interface(