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():
    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)

    model = LogisticRegression()
    model.fit(X_train, y_train)

    model_filename = "/mnt/data/model.pkl"
    dump(model, model_filename)
    
    return f"Modelo treinado e salvo em: {model_filename}"

iface = gr.Interface(fn=train_model, inputs=[], outputs=["text"])
iface.launch()