import gradio as gr
from sklearn.datasets import load_iris
from sklearn.model_selection import train_test_split
from sklearn.linear_model import LogisticRegression
from joblib import dump
import os

def train_model():
    # Carregar e dividir o dataset
    data = load_iris()
    X = data.data
    y = data.target
    X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=0.2, random_state=42)

    # Treinando o modelo
    model = LogisticRegression()
    model.fit(X_train, y_train)

    # Criar diretório para salvar o modelo
    diretorio = "/mnt/data"
    os.makedirs(diretorio, exist_ok=True)

    # Salvar o modelo em /mnt/data/
    model_filename = os.path.join(diretorio, "model.pkl")
    dump(model, model_filename)
    
    return f"Modelo treinado e salvo em: {model_filename}"

# Inicialize a interface Gradio
iface = gr.Interface(fn=train_model, inputs=[], outputs=["text"])

# Inicie a aplicação
iface.launch()