Fix

README.md

@@ -16,14 +16,14 @@ license: cc-by-nc-sa-4.0
 
 ```python
 import torch
-from transformers import AutoTokenizer, AutoModel
+from transformers import AutoTokenizer, AutoModelForSequenceClassification, pipeline
 
 tokenizer = AutoTokenizer.from_pretrained("Yehor/ual-topics-classifier")
-model = AutoModel.from_pretrained("Yehor/ual-topics-classifier")
+model = AutoModelForSequenceClassification.from_pretrained("Yehor/ual-topics-classifier")
 
 device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
 
-topic_classifier = pipeline(task='text-classification', model=model, tokenizer=tokenizer, device=device)
+topic_classifier = pipeline(task='text-classification', model=model, tokenizer=tokenizer, device=device, top_k=5)
 
 question = """
 Що мені робити на ВЛК
@@ -37,3 +37,11 @@ question = """
 
 print(topic_classifier(question))
 ```
+
+Results:
+
+```
+[[{'label': 'viiskovie_pravo', 'score': 0.9991637468338013}, {'label': 'inshe', 'score': 0.00035533116897568107}, {'label': 'immighratsiia_iemighratsiia', 'score': 8.094120130408555e-05}, {'label': 'pratsevlashtuvvannya', 'score': 4.537589484243654e-05}, {'label': 'opodatkuvannia', 'score': 3.591874337871559e-05}]]
+
+[[{'label': 'simejne_pravo', 'score': 0.9992771744728088}, {'label': 'sotsialnyj_zakhist', 'score': 0.0002545009192544967}, {'label': 'piensiiata_sotsialni_viplati', 'score': 5.3493567975237966e-05}, {'label': 'sudova_praktika', 'score': 3.8876780308783054e-05}, {'label': 'doghovirni_vidnosini', 'score': 3.169032061123289e-05}]]
+```