fix

app.py

@@ -5,39 +5,51 @@ from sklearn.ensemble import RandomForestRegressor
 from sklearn.metrics import r2_score
 import matplotlib.pyplot as plt
 
-# Загружаем датасет один раз
+# Загружаем датасет
 ds = load_dataset("QSBench/QSBench-Core-v1.0.0-demo")
 
-# Функция для отображения таблицы
+# Определяем доступные сплиты (предполагаем train, val, test)
+available_splits = list(ds.keys())
+print("Available splits:", available_splits)  # для отладки в логах
+
+# Функция для отображения данных выбранного сплита
 def show_data(split):
     try:
         df = pd.DataFrame(ds[split])
         return df.head(10)
     except Exception as e:
-        return f"Error loading data: {e}"
+        # Возвращаем пустой DataFrame с сообщением об ошибке
+        return pd.DataFrame({"Error": [str(e)]})
 
 # Функция для обучения модели и создания графика
 def train_model():
-    feature_cols = ["total_gates", "gate_entropy", "meyer_wallach"]
-    target_col = "ideal_expval_Z_global"
+    try:
+        # Проверяем, есть ли нужные сплиты
+        if "train" not in ds or "test" not in ds:
+            return None  # или вернуть сообщение
+        feature_cols = ["total_gates", "gate_entropy", "meyer_wallach"]
+        target_col = "ideal_expval_Z_global"
 
-    X_train = pd.DataFrame(ds["train"])[feature_cols]
-    y_train = pd.DataFrame(ds["train"])[target_col]
-    X_test = pd.DataFrame(ds["test"])[feature_cols]
-    y_test = pd.DataFrame(ds["test"])[target_col]
+        X_train = pd.DataFrame(ds["train"])[feature_cols]
+        y_train = pd.DataFrame(ds["train"])[target_col]
+        X_test = pd.DataFrame(ds["test"])[feature_cols]
+        y_test = pd.DataFrame(ds["test"])[target_col]
 
-    model = RandomForestRegressor(n_estimators=100, random_state=42)
-    model.fit(X_train, y_train)
-    y_pred = model.predict(X_test)
-    r2 = r2_score(y_test, y_pred)
+        model = RandomForestRegressor(n_estimators=100, random_state=42)
+        model.fit(X_train, y_train)
+        y_pred = model.predict(X_test)
+        r2 = r2_score(y_test, y_pred)
 
-    fig, ax = plt.subplots()
-    ax.scatter(y_test, y_pred, alpha=0.5)
-    ax.plot([y_test.min(), y_test.max()], [y_test.min(), y_test.max()], 'r--')
-    ax.set_xlabel("True value")
-    ax.set_ylabel("Predicted")
-    ax.set_title(f"Predictions vs. Truth (R² = {r2:.4f})")
-    return fig
+        fig, ax = plt.subplots()
+        ax.scatter(y_test, y_pred, alpha=0.5)
+        ax.plot([y_test.min(), y_test.max()], [y_test.min(), y_test.max()], 'r--')
+        ax.set_xlabel("True value")
+        ax.set_ylabel("Predicted")
+        ax.set_title(f"Predictions vs. Truth (R² = {r2:.4f})")
+        return fig
+    except Exception as e:
+        # В случае ошибки возвращаем сообщение как график? проще вернуть None и показать текст
+        raise e
 
 # Интерфейс
 with gr.Blocks(title="QSBench Demo Explorer") as demo:
@@ -52,15 +64,17 @@ with gr.Blocks(title="QSBench Demo Explorer") as demo:
 
     with gr.Tabs():
         with gr.TabItem("Data Explorer"):
+            # Определяем список сплитов из загруженного датасета
+            split_choices = available_splits if available_splits else ["train", "val", "test"]
             split_selector = gr.Dropdown(
-                choices=["train", "validation", "test"],
+                choices=split_choices,
                 label="Choose a split",
-                value="train"
+                value=split_choices[0] if split_choices else "train"
             )
             data_table = gr.Dataframe(label="First 10 rows", interactive=False)
             split_selector.change(fn=show_data, inputs=split_selector, outputs=data_table)
             # Загружаем данные по умолчанию
-            demo.load(fn=lambda: show_data("train"), outputs=data_table)
+            demo.load(fn=lambda: show_data(split_choices[0] if split_choices else "train"), outputs=data_table)
 
         with gr.TabItem("Model Demo"):
             train_button = gr.Button("Train Random Forest")