Update app.py

app.py

@@ -1,34 +1,39 @@
 import gradio as gr
-import datetime
-import torch
 import os
-from transformers import AutoTokenizer, AutoModelForSequenceClassification
-from datasets import Dataset, DatasetDict, disable_caching
+import json
+import uuid
+import torch
+import datetime
 import pandas as pd
-from huggingface_hub import HfApi, HfFolder
+from transformers import AutoTokenizer, AutoModelForSequenceClassification
+from huggingface_hub import HfApi, create_repo, upload_file
+from datasets import Dataset
 
-# CONFIG
-MODEL_NAME = "distilbert-base-uncased-finetuned-sst-2-english"  # Change if needed
-HF_DATASET_REPO = "M2ai/mgtd-logs"  # Must be created beforehand
-# Token from environment in Spaces
+# Configuration
+MODEL_NAME = "distilbert-base-uncased-finetuned-sst-2-english"
+HF_DATASET_REPO = "M2ai/mgtd-logs"
 HF_TOKEN = os.getenv("Mgtd")
-# Load model + tokenizer
+DATASET_CREATED = False
+
+# Load model and tokenizer
 tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)
 model = AutoModelForSequenceClassification.from_pretrained(MODEL_NAME)
-# Log entries
-log_entries = []
+
+# Make directories
+os.makedirs("logs", exist_ok=True)
 
 def setup_hf_dataset():
     global DATASET_CREATED
     if not DATASET_CREATED and HF_TOKEN:
         try:
             api = HfApi()
-            create_repo(DATASET_NAME, repo_type="dataset", token=HF_TOKEN, exist_ok=True)
+            create_repo(HF_DATASET_REPO, repo_type="dataset", token=HF_TOKEN, exist_ok=True)
             DATASET_CREATED = True
-            print(f"Dataset {DATASET_NAME} is ready")
-        except Exception as e: print(f"Error setting up dataset: {e}")
+            print(f"Dataset {HF_DATASET_REPO} is ready")
+        except Exception as e:
+            print(f"Error setting up dataset: {e}")
     elif not HF_TOKEN:
-        print("Warning: HF_TOKEN not set. Data will be stored locally only.")
+        print("Warning: HF_TOKEN not set. Logs will be saved locally only.")
 
 def infer_and_log(text_input):
     inputs = tokenizer(text_input, return_tensors="pt", truncation=True)
@@ -36,39 +41,49 @@ def infer_and_log(text_input):
         outputs = model(**inputs)
         logits = outputs.logits.tolist()
         predicted = torch.argmax(outputs.logits, dim=-1).item()
-        output_label = model.config.id2label[predicted]
+        label = model.config.id2label[predicted]
 
-    log_entries.append({
-        "timestamp": datetime.datetime.now().isoformat(),
+    timestamp = datetime.datetime.now().isoformat()
+    submission_id = str(uuid.uuid4())
+    log_data = {
+        "id": submission_id,
+        "timestamp": timestamp,
         "input": text_input,
-        "logits": logits,
-    })
+        "logits": logits
+    }
 
-    return output_label
+    log_file = f"logs/{timestamp.replace(':', '_')}.json"
+    with open(log_file, "w") as f:
+        json.dump(log_data, f, indent=2)
 
-def clear_fields():
-    return "", ""
+    if HF_TOKEN and DATASET_CREATED:
+        try:
+            api = HfApi()
+            api.upload_file(
+                path_or_fileobj=log_file,
+                path_in_repo=f"logs/{os.path.basename(log_file)}",
+                repo_id=HF_DATASET_REPO,
+                repo_type="dataset",
+                token=HF_TOKEN
+            )
+            print(f"Uploaded log {submission_id} to {HF_DATASET_REPO}")
+        except Exception as e:
+            print(f"Error uploading to HF dataset: {e}")
 
-def save_to_hf():
-    if not HF_TOKEN:
-        return "No Hugging Face token found in environment. Cannot push dataset."
+    return label
 
-    if not log_entries:
-        return "No logs to push."
+def clear_fields():
+    return "", ""
 
-    df = pd.DataFrame(log_entries)
-    dataset = Dataset.from_pandas(df)
-    
-    dataset.push_to_hub(HF_DATASET_REPO, token=HF_TOKEN)
-    log_entries.clear()
-    return f"Pushed {len(df)} logs to {HF_DATASET_REPO}!"
+# Setup the dataset on startup
+setup_hf_dataset()
 
-with gr.Blocks() as demo:
-    gr.Markdown("## AI-generated text detector")
+with gr.Blocks() as app:
+    gr.Markdown("## AI Text Detector")
 
     with gr.Row():
-        input_box = gr.Textbox(label="Input Text", lines=18, interactive=True)
-        output_box = gr.Textbox(label="Output Detected", lines=18)
+        input_box = gr.Textbox(label="Input Text", lines=10, interactive=True)
+        output_box = gr.Textbox(label="Output", lines=2, interactive=False)
 
     with gr.Row():
         submit_btn = gr.Button("Submit")
@@ -78,4 +93,4 @@ with gr.Blocks() as demo:
     clear_btn.click(fn=clear_fields, outputs=[input_box, output_box])
 
 if __name__ == "__main__":
-    demo.launch()
+    app.launch()