app.py

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):
    try:
        # Read the CSV file
        df = pd.read_csv(file.name)  # Use file.name to get the file path
        
        # 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)
        
        # Add predicted classes to the DataFrame
        df['predicted_class'] = predicted_classes.numpy()
        
        # Return the processed DataFrame as a CSV string
        output_csv = df.to_csv(index=False)
        
        return output_csv
    
    except Exception as e:
        return f"Error: {str(e)}"

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