model selection with speed ratings, time savings

app.py

@@ -8,16 +8,51 @@ import pickle
 import warnings
 warnings.filterwarnings("ignore", category=FutureWarning, module="transformers")
 
+# Model combinations with speed ratings and estimated time savings
+MODEL_COMBINATIONS = {
+    "Fastest (30 seconds)": {
+        "embedding": "sentence-transformers/all-MiniLM-L6-v2",
+        "generation": "distilgpt2",
+        "free": True,
+        "time_saved": "2.5 minutes"
+    },
+    "Balanced (1 minute)": {
+        "embedding": "sentence-transformers/all-MiniLM-L12-v2",
+        "generation": "facebook/opt-350m",
+        "free": True,
+        "time_saved": "2 minutes"
+    },
+    "High Quality (2 minutes)": {
+        "embedding": "sentence-transformers/all-mpnet-base-v2",
+        "generation": "gpt2",
+        "free": True,
+        "time_saved": "1 minute"
+    },
+    "Premium Speed (15 seconds)": {
+        "embedding": "sentence-transformers/all-MiniLM-L6-v2",
+        "generation": "microsoft/phi-1_5",
+        "free": False,
+        "time_saved": "2.75 minutes"
+    },
+    "Premium Quality (1.5 minutes)": {
+        "embedding": "openai-embedding-ada-002",
+        "generation": "meta-llama/Llama-2-7b-chat-hf",
+        "free": False,
+        "time_saved": "1.5 minutes"
+    }
+}
+
 @st.cache_resource
-def load_models():
+def load_models(model_combination):
     try:
-        tokenizer = AutoTokenizer.from_pretrained("distilbert-base-uncased")
-        embedding_model = AutoModel.from_pretrained("distilbert-base-uncased")
-        generation_model = AutoModelForCausalLM.from_pretrained("gpt2")
-        return tokenizer, embedding_model, generation_model
+        embedding_tokenizer = AutoTokenizer.from_pretrained(MODEL_COMBINATIONS[model_combination]['embedding'])
+        embedding_model = AutoModel.from_pretrained(MODEL_COMBINATIONS[model_combination]['embedding'])
+        generation_tokenizer = AutoTokenizer.from_pretrained(MODEL_COMBINATIONS[model_combination]['generation'])
+        generation_model = AutoModelForCausalLM.from_pretrained(MODEL_COMBINATIONS[model_combination]['generation'])
+        return embedding_tokenizer, embedding_model, generation_tokenizer, generation_model
     except Exception as e:
         st.error(f"Error loading models: {str(e)}")
-        return None, None, None
+        return None, None, None, None
 
 @st.cache_data
 def load_and_process_text(file_path):
@@ -95,18 +130,52 @@ st.markdown("""
     """, unsafe_allow_html=True)
 st.markdown('<p class="big-font">Chat with the Muse of A.R. Ammons. Ask questions or discuss poetry!</p>', unsafe_allow_html=True)
 
-# Load models and data
-with st.spinner("Loading models and data..."):
-    tokenizer, embedding_model, generation_model = load_models()
-    chunks, embeddings, index = load_data()
-    if chunks is None or embeddings is None or index is None:
+# Model selection
+if 'model_combination' not in st.session_state:
+    st.session_state.model_combination = "Fastest (30 seconds)"
+
+# Create a list of model options, with non-free models at the end
+free_models = [k for k, v in MODEL_COMBINATIONS.items() if v['free']]
+non_free_models = [k for k, v in MODEL_COMBINATIONS.items() if not v['free']]
+all_models = free_models + non_free_models
+
+# Custom CSS to grey out non-free options
+st.markdown("""
+<style>
+    .stSelectbox div[role="option"][aria-selected="false"]:nth-last-child(-n+2) {
+        color: grey !important;
+    }
+</style>
+""", unsafe_allow_html=True)
+
+selected_model = st.selectbox(
+    "Choose a model combination:",
+    all_models,
+    index=all_models.index(st.session_state.model_combination),
+    format_func=lambda x: f"{x} {'(Not Free)' if not MODEL_COMBINATIONS[x]['free'] else ''}"
+)
+
+# Prevent selection of non-free models
+if not MODEL_COMBINATIONS[selected_model]['free']:
+    st.warning("Premium models are not available in the free version.")
+    st.stop()
+
+st.session_state.model_combination = selected_model
+
+st.info(f"Potential time saved compared to slowest option: {MODEL_COMBINATIONS[selected_model]['time_saved']}")
+
+if st.button("Load Selected Models"):
+    with st.spinner("Loading models and data..."):
+        embedding_tokenizer, embedding_model, generation_tokenizer, generation_model = load_models(st.session_state.model_combination)
         chunks = load_and_process_text('ammons_muse.txt')
         embeddings = create_embeddings(chunks, embedding_model)
         index = create_faiss_index(embeddings)
-        save_data(chunks, embeddings, index)
+    
+    st.session_state.models_loaded = True
+    st.success("Models loaded successfully!")
 
-if tokenizer is None or embedding_model is None or generation_model is None or not chunks:
-    st.error("Failed to load necessary components. Please try again later.")
+if 'models_loaded' not in st.session_state or not st.session_state.models_loaded:
+    st.warning("Please load the models before chatting.")
     st.stop()
 
 # Initialize chat history