Update app.py

app.py

@@ -2,13 +2,13 @@ import streamlit as st
 from datasets import load_dataset
 from transformers import AutoTokenizer, AutoModelForSequenceClassification, pipeline
 import torch
+import pandas as pd
 
 def load_orca_dataset():
     st.info("Loading dataset... This may take a while.")
     return load_dataset("microsoft/orca-agentinstruct-1M-v1")
 
 @st.cache_data
-
 def load_model_and_tokenizer(model_name):
     tokenizer = AutoTokenizer.from_pretrained(model_name)
     model = AutoModelForSequenceClassification.from_pretrained(model_name)
@@ -40,11 +40,17 @@ def main():
     if "dataset" in st.session_state:
         dataset = st.session_state["dataset"]
 
+        # List available splits
+        available_splits = list(dataset.keys())
+        st.sidebar.subheader("Available Dataset Splits")
+        selected_split = st.sidebar.selectbox("Select Split", available_splits)
+
         st.subheader("Dataset Explorer")
-        st.write(dataset["train"].info)
+        st.write(f"Displaying information for split: `{selected_split}`")
+        st.write(dataset[selected_split].info)
 
         sample_size = st.slider("Number of Samples to Display", min_value=1, max_value=20, value=5)
-        st.write(dataset["train"].shuffle(seed=42).select(range(sample_size)))
+        st.write(dataset[selected_split].shuffle(seed=42).select(range(sample_size)))
 
         st.subheader("Model Evaluator")
         model_name = st.text_input("Enter Hugging Face Model Name", value="distilbert-base-uncased-finetuned-sst-2-english")
@@ -53,7 +59,7 @@ def main():
         if st.button("Load Model and Evaluate"):
             tokenizer, model = load_model_and_tokenizer(model_name)
 
-            results = evaluate_model(dataset["train"].shuffle(seed=42).select(range(max_samples)), tokenizer, model, max_samples)
+            results = evaluate_model(dataset[selected_split].shuffle(seed=42).select(range(max_samples)), tokenizer, model, max_samples)
 
             st.subheader("Evaluation Results")
             st.write(results)