import streamlit as st from datasets import load_dataset from transformer_ranker import TransformerRanker, prepare_popular_models st.title("Choose Your Transformer") model_options = { 'bert-tiny': 'prajjwal1/bert-tiny', 'bert-small': 'prajjwal1/bert-small', 'electra-small': 'google/electra-small-discriminator', 'deberta-small': 'microsoft/deberta-v3-small', 'distilbert-cased': 'distilbert-base-cased', 'distilbert-uncased': 'distilbert-base-uncased', } # 1) Select dataset names (from text classification or token classification subcategory) dataset_options = ['trec', 'conll2003'] # Example datasets; you can expand this list selected_dataset = st.selectbox("Select Dataset", dataset_options) # 2) Select the parameter for dataset downsampling downsample_values = [0.05, 0.1, 0.2, 0.3, 0.4, 0.5, 0.6, 0.7, 0.8, 0.9, 1.0] downsample_ratio = st.select_slider("Dataset Downsample Ratio", options=downsample_values, value=0.2) # 3) Select multiple models from HuggingFace model hub model_names = list(model_options.keys()) selected_models = st.multiselect("Select Models", model_names, default=['bert-tiny', 'electra-small']) selected_models = [model_options[model_name] for model_name in selected_models] # 4) Select the parameter for layer selection with layermean as the default layer_options = ['lastlayer', 'layermean', 'bestlayer'] selected_layer = st.selectbox("Layer Selection", layer_options, index=1) # Add real-time logging in the future log_expander = st.expander("Expand to view log") log_placeholder = log_expander.empty() # Placeholder for log updates # Button to run the ranking process if st.button("Run Model Ranking"): with st.spinner("Running the transformer-ranker..."): # Step 1: Load the selected dataset dataset = load_dataset(selected_dataset, trust_remote_code=True) # Step 2: Prepare the selected models language_models = prepare_popular_models('base') if selected_models == [] else selected_models # Step 3: Initialize the ranker ranker = TransformerRanker(dataset, dataset_downsample=downsample_ratio) # Placeholder for log updates log_placeholder.text("Real-time logging will be added here...") # Run the ranker results = ranker.run(language_models, batch_size=64) # Display the final results st.write(results) st.success("Ranking is Done!")