Update hugging_face.py

hugging_face.py

@@ -1,13 +1,15 @@
 import pandas as pd
+
 from datasets import load_dataset
 
-def load_data():
+def get_listings():
     dataset = load_dataset("MongoDB/airbnb_embeddings", streaming=True, split="train")
     dataset = dataset.take(100)
     dataset_df = pd.DataFrame(dataset)
     dataset_df.head(5)
     print("Columns:", dataset_df.columns)
     records = dataset_df.to_dict(orient='records')
+
     # To handle catch `NaT` values
     for record in records:
         for key, value in record.items():
@@ -19,10 +21,12 @@ def load_data():
             else:
                 if pd.isnull(value):
                     record[key] = None
+    
     try:
-      # Convert each dictionary to a Movie instance
-      listings = [Listing(**record).dict() for record in records]
-      # Get an overview of a single datapoint
-      print(listings[0].keys())
+        # Convert each dictionary to a Movie instance
+        listings = [Listing(**record).dict() for record in records]
+        # Get an overview of a single datapoint
+        print(listings[0].keys())
+        return listings
     except ValidationError as e:
-      print(e)
+        print(e)