import gradio as gr
import pandas as pd
from transformers import BertTokenizer, BertForSequenceClassification
import torch

# Load pre-trained TinyBERT model and tokenizer
tokenizer = BertTokenizer.from_pretrained('huawei-noah/TinyBERT_General_4L_312D')
model = BertForSequenceClassification.from_pretrained('huawei-noah/TinyBERT_General_4L_312D')

# Function to process the CSV file and generate predictions
def process_csv(file):
    # Read the CSV file
    df = pd.read_csv(file)
    
    # Ensure the CSV has a 'text' column
    if 'text' not in df.columns:
        return "Error: The CSV file must contain a 'text' column."
    
    # Tokenize the input text
    inputs = tokenizer(df['text'].tolist(), return_tensors='pt', padding=True, truncation=True)
    
    # Perform inference
    with torch.no_grad():
        outputs = model(**inputs)
    
    # Get predicted classes
    _, predicted_classes = torch.max(outputs.logits, dim=1)
    df['predicted_class'] = predicted_classes.numpy()
    
    # Return the processed DataFrame as a CSV string
    return df.to_csv(index=False)

# Create the Gradio interface
input_csv = gr.File(label="Upload CSV File")
output_csv = gr.File(label="Download Processed CSV")

demo = gr.Interface(
    fn=process_csv,
    inputs=input_csv,
    outputs=output_csv,
    title="CSV Data Processing with TinyBERT",
    description="Upload a CSV file with a 'text' column, and the model will process the data and provide predictions."
)

# Launch the Gradio interface
demo.launch()