Use right bert

app.py

@@ -10,9 +10,8 @@ def greet_json():
 from transformers import BertTokenizer, BertForSequenceClassification
 import torch
 
-from bertopic import BERTopic
 
-model = BERTopic.load("sleiyer/restricted_item_detector")
+model = BertForSequenceClassification.load("sleiyer/restricted_item_detector")
 # Load the trained model and tokenizer
 tokenizer = BertTokenizer.from_pretrained('bert-base-uncased')
 

restrictedItems/train.py

@@ -74,7 +74,6 @@ dataset = restricted_dataset + normal_dataset
 
 # Split dataset
 train_texts, val_texts, train_labels, val_labels = train_test_split([item[0] for item in dataset], [item[1] for item in dataset], test_size=0.2)
-import pdb; pdb.set_trace()
 
 # Load pre-trained BERT tokenizer
 tokenizer = BertTokenizer.from_pretrained('bert-base-uncased')
@@ -129,4 +128,4 @@ trainer.train()
 # Evaluate model
 trainer.evaluate()
 
-model.save_pretrained('trained_model')
+model.push_to_hub("sleiyer/restricted_item_detector")