import gradio as gr
import joblib
import pandas as pd
from huggingface_hub import hf_hub_download
import joblib
import os

hf_token = os.environ["hf_token"]

model_path = hf_hub_download(
    repo_id="kgemera/Biomedic_Text_Classifier",
    filename="classification_model.pkl",
    use_auth_token=hf_token
)

# Cargar el modelo usando la ruta absoluta
model = joblib.load(model_path)

# Single prediction
def predict_single(title, abstract):
    df = pd.DataFrame([{"title": title, "abstract": abstract}])
    return model.predict(df)[0]

single_interface = gr.Interface(
    fn=predict_single,
    inputs=[gr.Textbox(label="Title"), gr.Textbox(label="Abstract")],
    outputs=gr.Textbox(label="Predicted category")
)

# Batch prediction
def predict_batch(df):
    # df ya vendrá como dataframe desde gr.Dataframe
    return model.predict(df).tolist()

batch_interface = gr.Interface(
    fn=predict_batch,
    inputs=gr.Dataframe(
        headers=["title", "abstract"],
        label="Batch data",
        row_count=3
    ),
    outputs=gr.Textbox(label="Predictions")
)

# Lanzar ambos
single_interface.launch(server_name="0.0.0.0", server_port=7860, share=True)
batch_interface.launch(server_name="0.0.0.0", server_port=7861, share=True)