Update load_data.py

load_data.py

@@ -20,7 +20,7 @@ HUB_DATASET_NAME = os.environ.get('HUB_DATASET_NAME') #get dataset name
 )
 def save_validated_to_hub(records, ctx):
     if len(records) > 0:
-        ds = rg.DatasetForTextClassification(records=records).to_datasets()   
+        ds = rg.DatasetForTokenClassification(records=records).to_datasets()   
         if HF_TOKEN:
             print("Pushing the dataset")
             print(ds)
@@ -57,9 +57,9 @@ class LoadDatasets:
         dataset = dataset.remove_columns("metrics")
         records = rg.DatasetForTokenClassification.from_datasets(dataset)
 
-        #settings = rg.TokenClassificationSettings(
-        #    label_schema=["BAD INSTRUCTION", "BAD INPUT", "BAD OUTPUT", "INAPPROPRIATE", "BIASED", "ALL GOOD"]
-        #)
+        settings = rg.TokenClassificationSettings(
+            label_schema=["B-FUN", "I-FUN", "B-RES", "I-RES", "B-LEVEL", "I-LEVEL"]
+        )
         rg.configure_dataset(name="job-title-tagging", settings=settings, workspace="team")
         
         # Log the dataset