| import gradio as gr | |
| from transformers import AutoTokenizer, AutoModelForSequenceClassification | |
| tokenizer = AutoTokenizer.from_pretrained("WebOrganizer/TopicClassifier-NoURL") | |
| model = AutoModelForSequenceClassification.from_pretrained( | |
| "WebOrganizer/TopicClassifier-NoURL", | |
| trust_remote_code=True, | |
| use_memory_efficient_attention=False) | |
| def predict(text): | |
| inputs = tokenizer([text], return_tensors="pt") | |
| outputs = model(**inputs) | |
| probs = outputs.logits.softmax(dim=-1) | |
| pred_index = probs.argmax(dim=-1).item() | |
| confidence_score = probs[0, pred_index] | |
| id2label = model.config.id2label | |
| pred_label = id2label[pred_index] | |
| return {'topic': pred_label, 'confidence': confidence_score} | |
| title = "URL content Topic Categorizer" | |
| topic = gr.Interface( | |
| fn=predict, | |
| inputs='text', | |
| outputs='label', | |
| title=title, | |
| ) | |
| topic.launch() |