Update server.py

server.py

@@ -17,7 +17,7 @@ except ImportError:
     HF_DATASETS_AVAILABLE = False
     Features, Value = None, None
 
-STORAGE_BACKEND_CONFIG = os.getenv("STORAGE_BACKEND", "HF_DATASET").upper()
+STORAGE_BACKEND_CONFIG = os.getenv("STORAGE_BACKEND", "JSON").upper()
 HF_DATASET_REPO = os.getenv("HF_DATASET_REPO")
 HF_TOKEN = os.getenv("HF_TOKEN")
 DB_FILE_JSON = "social_data.json"
@@ -104,17 +104,29 @@ def load_data():
             if all([HF_DATASETS_AVAILABLE, HF_TOKEN, HF_DATASET_REPO]):
                 try:
                     ds_dict = load_dataset(HF_DATASET_REPO, token=HF_TOKEN, trust_remote_code=True)
-                    users = dict(zip(ds_dict['users']['username'], ds_dict['users']['password']))
-                    posts = ds_dict['posts'].to_pandas()
-                    comments = ds_dict['comments'].to_pandas()
+                    if ds_dict and all(k in ds_dict for k in ['users', 'posts', 'comments']):
+                        users = dict(zip(ds_dict['users']['username'], ds_dict['users']['password']))
+                        posts = ds_dict['posts'].to_pandas()
+                        comments = ds_dict['comments'].to_pandas()
+                        print("Successfully loaded data from HF Dataset.")
+                    else:
+                        raise ValueError("Dataset dictionary is empty or malformed.")
                 except Exception as e:
                     print(f"Could not load from HF Dataset '{HF_DATASET_REPO}'. Attempting to initialize. Error: {e}")
                     try:
                         user_features = Features({'username': Value('string'), 'password': Value('string')})
                         post_features = Features({'post_id': Value('int64'), 'username': Value('string'), 'content': Value('string'), 'timestamp': Value('string')})
                         comment_features = Features({'comment_id': Value('int64'), 'post_id': Value('int64'), 'username': Value('string'), 'content': Value('string'), 'timestamp': Value('string'), 'reply_to_comment_id': Value('int64')})
-                        dataset_dict = DatasetDict({'users': Dataset.from_pandas(pd.DataFrame(list(users.items()), columns=['username', 'password']), features=user_features), 'posts': Dataset.from_pandas(posts, features=post_features), 'comments': Dataset.from_pandas(comments, features=comment_features)})
+                        
+                        initial_users_df = pd.DataFrame(list(users.items()), columns=['username', 'password'])
+                        
+                        dataset_dict = DatasetDict({
+                            'users': Dataset.from_pandas(initial_users_df, features=user_features),
+                            'posts': Dataset.from_pandas(posts, features=post_features),
+                            'comments': Dataset.from_pandas(comments, features=comment_features)
+                        })
                         dataset_dict.push_to_hub(HF_DATASET_REPO, token=HF_TOKEN, private=True)
+                        print(f"Successfully initialized new empty HF Dataset at {HF_DATASET_REPO}.")
                     except Exception as e_push:
                         print(f"CRITICAL: Failed to create new HF Dataset. Falling back to RAM. Push Error: {e_push}")
                         STORAGE_BACKEND_CONFIG = "RAM"
@@ -243,4 +255,4 @@ with gr.Blocks(theme=gr.themes.Soft(), title="Social App") as demo:
         gr.Interface(api_get_feed, None, "dataframe", api_name="get_feed")
 
 if __name__ == "__main__":
-    demo.queue().launch(server_name="0.0.0.0", server_port=7860, share=False)
+    demo.queue().launch(server_name="0.0.0.0", server_port=7861, share=False)