Update app.py

app.py

@@ -1,24 +1,24 @@
 import gradio as gr
 import datetime
 import torch
+import os
 from transformers import AutoTokenizer, AutoModelForSequenceClassification
-from datasets import load_dataset, Dataset, DatasetDict
-import huggingface_hub
+from datasets import Dataset, DatasetDict, disable_caching
 import pandas as pd
+from huggingface_hub import HfApi, HfFolder
 
 # CONFIG
-MODEL_NAME = "distilbert-base-uncased-finetuned-sst-2-english"  # replace with your own
-HF_DATASET_REPO = "your-username/your-logging-dataset"  # create on HF Hub
-HF_TOKEN = "hf_..."  # your Hugging Face token with write access
+MODEL_NAME = "distilbert-base-uncased-finetuned-sst-2-english"  # Change if needed
+HF_DATASET_REPO = "your-username/your-logging-dataset"  # Must be created beforehand
+
+# Token from environment in Spaces
+HF_TOKEN = os.getenv("HUGGINGFACE_TOKEN")
 
 # Load model + tokenizer
 tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)
 model = AutoModelForSequenceClassification.from_pretrained(MODEL_NAME)
 
-# Setup dataset pushing
-huggingface_hub.login(token=HF_TOKEN)
-
-# Store session logs
+# Log entries
 log_entries = []
 
 def infer_and_log(text_input):
@@ -29,7 +29,6 @@ def infer_and_log(text_input):
         predicted = torch.argmax(outputs.logits, dim=-1).item()
         output_label = model.config.id2label[predicted]
 
-    # Create log entry
     log_entries.append({
         "timestamp": datetime.datetime.now().isoformat(),
         "input": text_input,
@@ -42,20 +41,25 @@ def clear_fields():
     return "", ""
 
 def save_to_hf():
+    if not HF_TOKEN:
+        return "No Hugging Face token found in environment. Cannot push dataset."
+
     if not log_entries:
-        return "Nothing to save."
+        return "No logs to push."
+
+    df = pd.DataFrame(log_entries)
+    dataset = Dataset.from_pandas(df)
     
-    dataset = Dataset.from_pandas(pd.DataFrame(log_entries))
-    dataset.push_to_hub(HF_DATASET_REPO)
+    dataset.push_to_hub(HF_DATASET_REPO, token=HF_TOKEN)
     log_entries.clear()
-    return "Data saved to Hugging Face!"
+    return f"Pushed {len(df)} logs to {HF_DATASET_REPO}!"
 
 with gr.Blocks() as demo:
-    gr.Markdown("### 🔤 Text Classification Demo")
+    gr.Markdown("## 🤖 Text Classification with Logging")
 
     with gr.Row():
-        input_box = gr.Textbox(label="Input Text", lines=5, interactive=True)
-        output_box = gr.Textbox(label="Predicted Label", lines=5)
+        input_box = gr.Textbox(label="Input Text", lines=4, interactive=True)
+        output_box = gr.Textbox(label="Predicted Label", lines=2)
 
     with gr.Row():
         submit_btn = gr.Button("Submit")
@@ -66,7 +70,7 @@ with gr.Blocks() as demo:
     submit_btn.click(fn=infer_and_log, inputs=input_box, outputs=output_box)
     clear_btn.click(fn=clear_fields, outputs=[input_box, output_box])
     
-    gr.Button("Save Logs to Hub").click(fn=save_to_hf, outputs=status_box)
+    gr.Button("Save Logs to HF Dataset").click(fn=save_to_hf, outputs=status_box)
 
 if __name__ == "__main__":
     demo.launch()