Update finetune3.py

finetune3.py

@@ -61,6 +61,9 @@ def prepare_data(df, tokenizer):
 def main():
     st.title("Patent Classification with Fine-Tuned BERT")
     
+    # Initialize model directory path
+    model_dir = './finetuned_model'
+    
     # Load data
     df = load_data()
     
@@ -70,12 +73,11 @@ def main():
     
     # Prepare data
     model_name = "bert-base-uncased"
-    dummy_num_labels = 5
-    tokenizer, model = load_tokenizer_and_model(model_name, dummy_num_labels)
+    tokenizer, model = load_tokenizer_and_model(model_name, num_labels=5)
     dataset, num_labels = prepare_data(df, tokenizer)
     
     # Update the model with the correct number of labels based on the data
-    if num_labels != dummy_num_labels:
+    if num_labels != 5:
         model = AutoModelForSequenceClassification.from_pretrained(model_name, num_labels=num_labels)
     
     # Split the dataset
@@ -102,7 +104,7 @@ def main():
 
     # Fine-tune model
     training_args = TrainingArguments(
-        output_dir='./results',
+        output_dir=model_dir,
         evaluation_strategy="epoch",
         learning_rate=2e-5,
         per_device_train_batch_size=8,
@@ -123,15 +125,13 @@ def main():
     if st.button('Train Model'):
         with st.spinner('Training in progress...'):
             trainer.train()
-            model.save_pretrained("./finetuned_model")
-            tokenizer.save_pretrained("./finetuned_model")
+            model.save_pretrained(model_dir)
+            tokenizer.save_pretrained(model_dir)
             st.success("Model training complete and saved.")
     
     # Display pretrained model data
     st.subheader("Pretrained Model")
     if st.button('Show Pretrained Model'):
-        model_dir = './finetuned_model'
-        
         # List files in the directory
         if os.path.exists(model_dir):
             files = os.listdir(model_dir)