# Import necessary libraries
import torch
import gradio as gr
from transformers import AutoModelForSequenceClassification, AutoTokenizer
from train import train_model

# Load pre-trained model and tokenizer
model_name = "bert-base-uncased"
model = AutoModelForSequenceClassification.from_pretrained(model_name)
tokenizer = AutoTokenizer.from_pretrained(model_name)

# Train the model
train_model(model, tokenizer)

# Create a Gradio interface
interface = gr.Interface(
    fn=model,
    inputs="text",
    outputs="text",
    title="Hugging Face Chatbot",
    description="A Gradio chatbot space using Hugging Face inference"
)

# Launch the interface
interface.launch()