feat: memoize model

app.py

@@ -36,7 +36,7 @@ def display_streamlit_sidebar():
     do_sample = form.checkbox('do_sample', value=st.session_state["do_sample"])
     top_p = form.slider(label="top_p", min_value=0.0, max_value=1.0, step=0.01, value=st.session_state["top_p"])
     top_k = form.slider(label="top_k", min_value=1, max_value=1000, step=1, value=st.session_state["top_k"])
-    max_new_tokens = form.slider(label="max_new_tokens", min_value=32, max_value=16384, step=1, value=st.session_state["max_new_tokens"])
+    max_new_tokens = form.slider(label="max_new_tokens", min_value=32, max_value=4096, step=1, value=st.session_state["max_new_tokens"])
     repetition_penalty = form.slider(label="repetition_penalty", min_value=0.0, max_value=5.0, step=0.01, value=st.session_state["repetition_penalty"])
 
     submitted = form.form_submit_button("Start session")
@@ -76,7 +76,7 @@ def init_session_state():
     st.session_state["do_sample"] = True
     st.session_state["top_p"] = 0.95
     st.session_state["top_k"] = 40
-    st.session_state["max_new_tokens"] = 512
+    st.session_state["max_new_tokens"] = 4096
     st.session_state["repetition_penalty"] = 1.1
     st.session_state["system_prompt"] = "You are a medical expert that provides answers for a medically trained audience"
     st.session_state["prompt"] = ""
@@ -143,9 +143,9 @@ def display_llm_output():
     form = st.form('llm')
 
     prompt_format_str = get_prompt_format(st.session_state["model_name_or_path"])
-    prompt_format = form.text_area('Prompt format', value=prompt_format_str)
-    system_prompt = form.text_area('System prompt', value=st.session_state["system_prompt"])
-    prompt = form.text_area('Prompt', value=st.session_state["prompt"])
+    prompt_format = form.text_area('Prompt format', value=prompt_format_str, height=300)
+    system_prompt = form.text_area('System message', value=st.session_state["system_prompt"], height=300)
+    prompt = form.text_area('Prompt', value=st.session_state["prompt"], height=400)
 
     submitted = form.form_submit_button('Submit')
 

utils/epfl_meditron_utils.py

@@ -9,11 +9,19 @@ def gptq_model_options():
         "TheBloke/meditron-70B-GPTQ",
     ]
 
+loaded_model = None
+loaded_model_name = ""
+
 def get_llm_response(model_name_or_path, temperature, do_sample, top_p, top_k, max_new_tokens, repetition_penalty, formatted_prompt):
-    model = AutoModelForCausalLM.from_pretrained(model_name_or_path,
+    if loaded_model != model_name_or_path:
+        global loaded_model
+        global loaded_model_name
+        loaded_model = AutoModelForCausalLM.from_pretrained(model_name_or_path,
                                                 device_map="auto",
                                                 trust_remote_code=False,
                                                 revision="main")
+        loaded_model_name = model_name_or_path
+        
 
     tokenizer = AutoTokenizer.from_pretrained(model_name_or_path, use_fast=True)
             
@@ -30,7 +38,7 @@ def get_llm_response(model_name_or_path, temperature, do_sample, top_p, top_k, m
     print("*** Pipeline:")
     pipe = pipeline(
         "text-generation",
-        model=model,
+        model=loaded_model,
         tokenizer=tokenizer,
         max_new_tokens=max_new_tokens,
         do_sample=do_sample,