Update finetune1.py

taking 5 entries only

finetune1.py

@@ -5,6 +5,7 @@ from datasets import Dataset
 from sklearn.model_selection import train_test_split
 import requests
 from io import BytesIO
+import numpy as np
 
 # Load the dataset
 @st.cache_data
@@ -17,9 +18,9 @@ def load_data():
 
 # 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, num_labels=2)  # Adjust num_labels as needed
+    model = AutoModelForSequenceClassification.from_pretrained(model_name, num_labels=num_labels)
     return tokenizer, model
 
 # Tokenize and prepare the dataset
@@ -27,10 +28,17 @@ 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()
+    # 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)]
     
-    # Assuming `patent_number` is categorical or needs to be converted to labels
-    labels = jan_2016_df['patent_number'].astype('category').cat.codes.tolist()
+    # 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):
@@ -48,7 +56,7 @@ def prepare_data(df, tokenizer):
     
     dataset = Dataset.from_dict(dataset_dict)
     
-    return dataset
+    return dataset, num_labels
 
 # Define Streamlit app
 def main():
@@ -61,12 +69,13 @@ def main():
     st.subheader("Sample Data from January 2016")
     st.write(df.head())
     
-    # Load tokenizer and model
-    model_name = "bert-base-uncased"
-    tokenizer, model = load_tokenizer_and_model(model_name)
-    
     # Prepare data
-    dataset = prepare_data(df, tokenizer)
+    tokenizer, model = None, None
+    dataset, num_labels = prepare_data(df, tokenizer)
+    
+    # Load tokenizer and model with the correct number of labels
+    model_name = "bert-base-uncased"
+    tokenizer, model = load_tokenizer_and_model(model_name, 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)
@@ -114,4 +123,4 @@ def main():
         st.write("Pretrained model is `bert-base-uncased`. Fine-tuned model is saved at './finetuned_model'.")
 
 if __name__ == "__main__":
-    main()
+    main()