import gradio as gr
import datetime
import torch
from transformers import AutoTokenizer, AutoModelForSequenceClassification
from datasets import load_dataset, Dataset, DatasetDict
import huggingface_hub
import pandas as pd

# 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

# 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 = []

def infer_and_log(text_input):
    inputs = tokenizer(text_input, return_tensors="pt", truncation=True)
    with torch.no_grad():
        outputs = model(**inputs)
        logits = outputs.logits.tolist()
        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,
        "logits": logits,
    })

    return output_label

def clear_fields():
    return "", ""

def save_to_hf():
    if not log_entries:
        return "Nothing to save."
    
    dataset = Dataset.from_pandas(pd.DataFrame(log_entries))
    dataset.push_to_hub(HF_DATASET_REPO)
    log_entries.clear()
    return "Data saved to Hugging Face!"

with gr.Blocks() as demo:
    gr.Markdown("### 🔤 Text Classification Demo")

    with gr.Row():
        input_box = gr.Textbox(label="Input Text", lines=5, interactive=True)
        output_box = gr.Textbox(label="Predicted Label", lines=5)

    with gr.Row():
        submit_btn = gr.Button("Submit")
        clear_btn = gr.Button("Clear")

    status_box = gr.Textbox(label="Status", interactive=False)

    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)

if __name__ == "__main__":
    demo.launch()