import gradio as gr
import pandas as pd
from transformers import TabularTransformerForSequenceClassification, TabularTransformerConfig
from transformers import Trainer, TrainingArguments
from datasets import Dataset
import torch

# Sample Data
data = {
    'feature1': [0.5, 0.3, 0.7, 0.2],
    'feature2': [1, 0, 1, 1],
    'feature3': [0.6, 0.1, 0.8, 0.4],
    'label': [0, 1, 0, 1]  # Binary classification
}
df = pd.DataFrame(data)
dataset = Dataset.from_pandas(df)

# Configure the Model
config = TabularTransformerConfig(
    num_labels=2,  # Binary classification
    numerical_features=['feature1', 'feature2', 'feature3']
)
model = TabularTransformerForSequenceClassification(config)

# Define Training Arguments
training_args = TrainingArguments(
    output_dir="./results",
    evaluation_strategy="epoch",
    learning_rate=2e-5,
    per_device_train_batch_size=4,
    num_train_epochs=3
)

# Define Trainer
trainer = Trainer(
    model=model,
    args=training_args,
    train_dataset=dataset,
    eval_dataset=dataset
)

# Train the model
trainer.train()

# Define Inference Function
def classify(feature1, feature2, feature3):
    input_data = {'feature1': feature1, 'feature2': feature2, 'feature3': feature3}
    input_df = pd.DataFrame([input_data])
    test_dataset = Dataset.from_pandas(input_df)
    with torch.no_grad():
        logits = model(**test_dataset[:][0]).logits
        prediction = torch.argmax(logits, dim=1).item()
        return "Class 1" if prediction == 1 else "Class 0"

# Gradio Interface
iface = gr.Interface(
    fn=classify,
    inputs=[
        gr.inputs.Slider(0, 1, step=0.1, label="Feature 1"),
        gr.inputs.Slider(0, 1, step=0.1, label="Feature 2"),
        gr.inputs.Slider(0, 1, step=0.1, label="Feature 3")
    ],
    outputs="text",
    title="Tabular Classification with Hugging Face",
    description="Classify entries based on tabular data"
)

iface.launch()