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

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

# Load dataset from Hugging Face repository
# Replace 'your-username' and 'your-dataset-name' with actual values
dataset = load_dataset('SharmaAmit1818/data_analysis/blob/main', data_files='data-qQeu1Z0CfsuqRUaDagRA1 (1).csv')

# Function to process the CSV file and generate predictions
def process_csv(file):
    try:
        # Read the CSV file using Pandas directly from the uploaded file object
        df = pd.read_csv(file)  # Use the file object directly
        
        # Debugging: Print the DataFrame shape and columns
        print(f"DataFrame shape: {df.shape}")
        print(f"DataFrame columns: {df.columns.tolist()}")

        # Check for 'text' column
        if 'text' not in df.columns:
            return "Error: The CSV file must contain a 'text' column."
        
        # Tokenize 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 predictions to DataFrame
        df['predicted_class'] = predicted_classes.numpy()
        
        # Return processed DataFrame as CSV string
        return df.to_csv(index=False)
    
    except FileNotFoundError:
        return "Error: The specified file was not found. Please check your upload."
    except pd.errors.EmptyDataError:
        return "Error: The uploaded file is empty."
    except pd.errors.ParserError:
        return "Error: There was an issue parsing the CSV file."
    except Exception as e:
        return f"An unexpected error occurred: {str(e)}"

# Create 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 Gradio interface
demo.launch()