File size: 2,204 Bytes
49b1ef8 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 |
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()
|