# Importer les bibliothèques nécessaires
import gradio as gr
from transformers import NVLM_D, 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)

# Fonction pour effectuer des prédictions
def predict(text):
    # Tokeniser le texte
    inputs = tokenizer(text, return_tensors="pt")
    
    # Effectuer la prédiction
    outputs = model(**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

# Créer l'interface Gradio
iface = gr.Interface(
    fn=predict,
    inputs="text",
    outputs="text",
    title="Prédiction avec NVLM-D",
    description="Entrez un texte pour obtenir une prédiction avec le modèle NVLM-D."
)

# Lancer l'application Gradio
iface.launch()