Update app/main.py

app/main.py

@@ -14,6 +14,7 @@ MODELS_PATH = "kadabengaran/IndoBERT-BiLSTM-Useful-App-Review"
 
 id2label= {0: 'Other', 1: 'Problem Discovery', 2: 'Information Seeking', 3: 'Feature Request'}
 label2id= {'Other': 0, 'Problem Discovery': 1, 'Information Seeking': 2, 'Feature Request': 3}
+numLabels= 4
 
 def get_device():
     if torch.cuda.is_available():
@@ -53,7 +54,7 @@ def remove_special_characters(text):
 def load_model():
     config = PeftConfig.from_pretrained(MODELS_PATH)
     inference_model = AutoModelForSequenceClassification.from_pretrained(
-        config.base_model_name_or_path, num_labels=2, id2label=id2label, label2id=label2id
+        config.base_model_name_or_path, num_labels=numLabels, id2label=id2label, label2id=label2id
     )
     tokenizer = AutoTokenizer.from_pretrained(config.base_model_name_or_path)
     model = PeftModel.from_pretrained(inference_model, MODELS_PATH)