Update finetune1.py

finetune1.py

@@ -1,5 +1,6 @@
-import streamlit as st
+import os
 import pandas as pd
+import streamlit as st
 from transformers import AutoTokenizer, AutoModelForSequenceClassification, Trainer, TrainingArguments
 from datasets import Dataset
 from sklearn.model_selection import train_test_split
@@ -60,26 +61,23 @@ 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()
     
-    # Show sample data
-    st.subheader("Sample Data from January 2016")
+    # Show data
+    st.subheader("Data from January 2016")
     st.write(df.head())
     
     # Prepare data
-    # First, select a model name
     model_name = "bert-base-uncased"
-    
-    # Initialize the tokenizer and model with a dummy number of labels for now
-    dummy_num_labels = 5
-    tokenizer, model = load_tokenizer_and_model(model_name, dummy_num_labels)
-    
-    # Prepare the data
+    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
@@ -95,9 +93,18 @@ def main():
     train_dataset = create_dataset(train_data)
     eval_dataset = create_dataset(eval_data)
     
+    # Show training data
+    st.subheader("Training Data")
+    train_df = pd.DataFrame({
+        'input_ids': [ids[:10] for ids in train_dataset['input_ids'][:5]], 
+        'attention_mask': [mask[:10] for mask in train_dataset['attention_mask'][:5]],
+        'labels': train_dataset['labels'][:5]
+    })
+    st.write(train_df)
+
     # 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,
@@ -118,14 +125,31 @@ 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'):
-        st.write("Pretrained model is `bert-base-uncased`. Fine-tuned model is saved at './finetuned_model'.")
+        if os.path.exists(model_dir):
+            # Show model name
+            st.write(f"Model name: `{model_name}`")
+            
+            # List .json files
+            json_files = [f for f in os.listdir(model_dir) if f.endswith('.json')]
+            if json_files:
+                st.write("Available `.json` files:")
+                for file in json_files:
+                    file_path = os.path.join(model_dir, file)
+                    with open(file_path, 'r', encoding='utf-8') as f:
+                        file_content = f.read()
+                    st.write(f"[{file}](data:file/{file})")  # Create clickable link
+                    st.text(file_content)  # Display file content
+            else:
+                st.write("No `.json` files found in `./finetuned_model` directory.")
+        else:
+            st.write("Directory `./finetuned_model` does not exist.")
 
 if __name__ == "__main__":
     main()