Updating column names

src/build_nomic.py

@@ -10,11 +10,16 @@ NOMIC_KEY = os.getenv('NOMIC_KEY')
 nomic.login(NOMIC_KEY)
 
 
+def count_words(text):
+    words = text.split()
+    return len(words)
+
+
 def build_nomic(dataset):
     df = dataset['train'].to_pandas()
 
-    non_embedding_columns = ['date_utc', 'title', 'flair', 'content', 'poster', 'permalink', 'id', 'content_length',
-                             'score', 'percentile_ranges']
+    non_embedding_columns = ['date_utc', 'title', 'flair', 'content', 'poster', 'permalink', 'id', 'word_count',
+                             'score', 'score_percentile']
 
     # Calculate the 0th, 10th, 20th, ..., 90th percentiles for the 'score' column
     percentiles = df['score'].quantile([0, .1, .2, .3, .4, .5, .6, .7, .8, .9]).tolist()
@@ -30,9 +35,11 @@ def build_nomic(dataset):
     # This assigns each score to its corresponding percentile range
     df['score_percentile'] = pd.cut(df['score'], bins=bins, labels=labels, include_lowest=True)
 
+    df['word_count'] = df['content'].apply(count_words)
+
     # Create Atlas project
     project = atlas.map_data(embeddings=np.stack(df['embedding'].values),
                              data=df[non_embedding_columns].to_dict(orient='records'),
                              id_field='id',
                              identifier='BORU Subreddit Neural Search',
-                             )
+                             )