Update app.py

app.py

@@ -7,25 +7,23 @@ from service_dops_api.dops_config import ServiceDopsConfig
 from service_dops_api.dops_classifier import DopsClassifier
 import json
 HF_TOKEN = os.getenv('HF_TOKEN')
-tokenizer_cat = BertTokenizer.from_pretrained("warleagle/service_name_categorizer",
+tokenizer_cat = BertTokenizer.from_pretrained("etadevosyan/service_categorizer_v2",
                                           token=HF_TOKEN)
-model_cat = BertForSequenceClassification.from_pretrained('warleagle/service_name_categorizer',token=HF_TOKEN)
+model_cat = BertForSequenceClassification.from_pretrained('etadevosyan/service_categorizer_v2',token=HF_TOKEN)
 
 clf_cat = pipeline("text-classification", model=model_cat, tokenizer=tokenizer_cat)
 
-id2label_cat = pd.read_pickle('id2label_service_categoriser.pickle')
 
-tokenizer_spec = BertTokenizer.from_pretrained("warleagle/specialists_categorizer_model",
+tokenizer_spec = BertTokenizer.from_pretrained("etadevosyan/specialists_categorizer_model",
                                           token=HF_TOKEN)
-model_spec = BertForSequenceClassification.from_pretrained('warleagle/specialists_categorizer_model',token=HF_TOKEN)
+model_spec = BertForSequenceClassification.from_pretrained('etadevosyan/specialists_categorizer_model',token=HF_TOKEN)
 
 clf_spec = pipeline("text-classification", model=model_spec, tokenizer=tokenizer_spec)
 
 id2label_spec = pd.read_pickle('id2label_spec_categoriser.pickle')
 def categoriser_predict(input_text):
     predictions = clf_cat(input_text)
-    numeric_label = int(predictions[0]['label'].split("_")[1])
-    text_label = id2label_cat[numeric_label]
+    text_label = predictions[0]['label']
     return text_label
 def doctor_spec_predict(input_text):
     predictions = clf_spec(input_text)