Move client instantiation

src/utilities.py

@@ -12,7 +12,6 @@ USERNAME = os.environ["USERNAME"]
 OG_DATASET = f"{USERNAME}/dataset-creator-reddit-{SUBREDDIT}"
 PROCESSED_DATASET = os.environ['PROCESSED_DATASET']
 
-client = Client("derek-thomas/nomic-embeddings")
 logger = setup_logger(__name__)
 
 
@@ -29,6 +28,9 @@ async def load_datasets():
 
 
 def merge_and_update_datasets(dataset, original_dataset):
+    # Get client
+    client = Client("derek-thomas/nomic-embeddings")
+
     # Merge and figure out which rows need to be updated with embeddings
     odf = original_dataset['train'].to_pandas()
     df = dataset['train'].to_pandas()
@@ -50,13 +52,13 @@ def merge_and_update_datasets(dataset, original_dataset):
     # Iterate over the DataFrame rows where 'embedding' is None
     for index, row in merged_df[merged_df['embedding'].isnull()].iterrows():
         # Update 'embedding' for the current row using our function
-        merged_df.at[index, 'embedding'] = update_embeddings(row['content'])
+        merged_df.at[index, 'embedding'] = update_embeddings(content=row['content'], client=client)
 
     dataset['train'] = Dataset.from_pandas(merged_df)
     logger.info(f"Updated {updated_rows} rows")
     return dataset
 
 
-def update_embeddings(content):
+def update_embeddings(content, client):
     embedding = client.predict(content, api_name="/embed")
     return np.array(embedding)