Create app.py

app.py

@@ -0,0 +1,46 @@
+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()