# Importer les bibliothèques nécessaires
import gradio as gr
from transformers import AutoModelForCausalLM, AutoTokenizer

# Charger le modèle et le tokenizer
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):
    # Tokeniser le texte
    inputs = tokenizer(text, return_tensors="pt")
    
    # Effectuer la prédiction
    outputs = model.generate(**inputs)
    
    # 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(
    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()