Update finetune3.py

finetune3.py

@@ -1,7 +1,8 @@
+import os
 import streamlit as st
 import pandas as pd
 from transformers import AutoTokenizer, AutoModelForSequenceClassification, Trainer, TrainingArguments
-from datasets import Dataset, DatasetDict
+from datasets import Dataset
 from sklearn.model_selection import train_test_split
 import requests
 from io import BytesIO
@@ -16,10 +17,9 @@ def load_data():
     return df
 
 # Tokenizer and model loading
-@st.cache_resource
-def load_tokenizer_and_model(model_name):
+def load_tokenizer_and_model(model_name, num_labels):
     tokenizer = AutoTokenizer.from_pretrained(model_name)
-    model = AutoModelForSequenceClassification.from_pretrained(model_name)
+    model = AutoModelForSequenceClassification.from_pretrained(model_name, num_labels=num_labels)
     return tokenizer, model
 
 # Tokenize and prepare the dataset
@@ -27,14 +27,26 @@ def prepare_data(df, tokenizer):
     df['filing_date'] = pd.to_datetime(df['filing_date'])
     jan_2016_df = df[df['filing_date'].dt.to_period('M') == '2016-01']
     
-    texts = jan_2016_df['invention_title'].tolist()
-    labels = jan_2016_df['patent_number'].tolist()
+    # Get only 5 unique labels
+    unique_labels = jan_2016_df['patent_number'].astype('category').cat.categories[:5]
+    jan_2016_df = jan_2016_df[jan_2016_df['patent_number'].isin(unique_labels)]
     
+    # Re-map labels to integers starting from 0
+    label_mapping = {label: idx for idx, label in enumerate(unique_labels)}
+    jan_2016_df['label'] = jan_2016_df['patent_number'].map(label_mapping)
+
+    texts = jan_2016_df['invention_title'].tolist()
+    labels = jan_2016_df['label'].tolist()
+    num_labels = len(unique_labels)
+
+    # Define tokenization function
     def tokenize_function(texts):
-        return tokenizer(texts, padding="max_length", truncation=True, return_tensors="pt", max_length=512)
-    
+        return tokenizer(texts, padding=True, truncation=True, return_tensors="pt", max_length=512)
+
+    # Tokenize texts
     tokenized_data = tokenize_function(texts)
-    
+
+    # Create dataset
     dataset_dict = {
         'input_ids': [x.tolist() for x in tokenized_data['input_ids']],
         'attention_mask': [x.tolist() for x in tokenized_data['attention_mask']],
@@ -43,16 +55,7 @@ def prepare_data(df, tokenizer):
     
     dataset = Dataset.from_dict(dataset_dict)
     
-    return dataset
-
-# Define a custom compute_loss function
-def compute_loss(model, inputs):
-    labels = inputs.get("labels")
-    outputs = model(**inputs)
-    logits = outputs.logits
-    loss_fct = torch.nn.CrossEntropyLoss()
-    loss = loss_fct(logits, labels)
-    return loss
+    return dataset, num_labels
 
 # Define Streamlit app
 def main():
@@ -65,27 +68,33 @@ def main():
     st.subheader("Sample Data from January 2016")
     st.write(df.head())
     
-    # Load tokenizer and model
+    # Prepare data
+    # First, select a model name
     model_name = "bert-base-uncased"
-    tokenizer, model = load_tokenizer_and_model(model_name)
     
-    # Prepare data
-    dataset = prepare_data(df, tokenizer)
+    # 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
+    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:
+        model = AutoModelForSequenceClassification.from_pretrained(model_name, num_labels=num_labels)
     
     # Split the dataset
     train_data, eval_data = train_test_split(list(zip(dataset['input_ids'], dataset['attention_mask'], dataset['labels'])), test_size=0.2, random_state=42)
     
-    train_dataset = Dataset.from_dict({
-        'input_ids': [item[0] for item in train_data],
-        'attention_mask': [item[1] for item in train_data],
-        'labels': [item[2] for item in train_data]
-    })
-    
-    eval_dataset = Dataset.from_dict({
-        'input_ids': [item[0] for item in eval_data],
-        'attention_mask': [item[1] for item in eval_data],
-        'labels': [item[2] for item in eval_data]
-    })
+    def create_dataset(data):
+        return Dataset.from_dict({
+            'input_ids': [item[0] for item in data],
+            'attention_mask': [item[1] for item in data],
+            'labels': [item[2] for item in data]
+        })
+
+    train_dataset = create_dataset(train_data)
+    eval_dataset = create_dataset(eval_data)
     
     # Fine-tune model
     training_args = TrainingArguments(
@@ -103,7 +112,7 @@ def main():
         args=training_args,
         train_dataset=train_dataset,
         eval_dataset=eval_dataset,
-        compute_loss=compute_loss  # Use the custom loss function
+        tokenizer=tokenizer
     )
     
     st.subheader("Training the Model")
@@ -117,7 +126,28 @@ def main():
     # 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'.")
+        model_dir = './finetuned_model'
+        
+        # List files in the directory
+        if os.path.exists(model_dir):
+            files = os.listdir(model_dir)
+            st.write("Contents of `./finetuned_model` directory:")
+            st.write(files)
+        else:
+            st.write("Directory `./finetuned_model` does not exist.")
+    
+    # Optionally, show the file content of model files (e.g., config.json, pytorch_model.bin)
+    st.subheader("Model File Contents")
+    if st.button('Show Model File Contents'):
+        if os.path.exists(model_dir):
+            files = os.listdir(model_dir)
+            for file in files:
+                file_path = os.path.join(model_dir, file)
+                with open(file_path, 'r') as f:
+                    st.write(f"**{file}:**")
+                    st.write(f.read())
+        else:
+            st.write("Directory `./finetuned_model` does not exist.")
 
 if __name__ == "__main__":
     main()