Update app.py

app.py

@@ -1,7 +1,7 @@
 from transformers import RobertaForSequenceClassification, RobertaTokenizer, Trainer, TrainingArguments
-from datasets import Dataset
+from datasets import Dataset, DatasetDict
 import pandas as pd
-import torch
+from sklearn.preprocessing import LabelEncoder
 
 # Load the dataset
 df = pd.read_csv("processed_step3.csv")
@@ -10,13 +10,20 @@ df = pd.read_csv("processed_step3.csv")
 def preprocess_data(row):
     return {"text": row["full_text"], "labels": row["narratives"]}
 
+# Apply label encoding to narratives to turn them into numeric labels
+label_encoder = LabelEncoder()
+df["labels"] = label_encoder.fit_transform(df["narratives"])
+
 # Create a Dataset object
-hf_dataset = Dataset.from_pandas(df).map(preprocess_data)
+hf_dataset = Dataset.from_pandas(df)
+
+# Split the dataset into train and validation sets (80/20 split)
+hf_dataset = hf_dataset.train_test_split(test_size=0.2)
 
 # Load pre-trained tokenizer and model
 tokenizer = RobertaTokenizer.from_pretrained("roberta-base")
 model = RobertaForSequenceClassification.from_pretrained(
-    "roberta-base", num_labels=len(set(df["narratives"])))
+    "roberta-base", num_labels=len(label_encoder.classes_))  # Use the number of unique labels
 
 # Tokenize the data
 def tokenize_function(examples):
@@ -42,8 +49,8 @@ training_args = TrainingArguments(
 trainer = Trainer(
     model=model,
     args=training_args,
-    train_dataset=hf_dataset["train"],
-    eval_dataset=hf_dataset["validation"],
+    train_dataset=hf_dataset["train"],  # Train set
+    eval_dataset=hf_dataset["test"],  # Validation set
     tokenizer=tokenizer
 )