feat: specify params

app.py

@@ -11,7 +11,7 @@ import sys
 import io
 
 from utils.default_values import get_system_prompt, get_guidelines_dict
-from utils.epfl_meditron_utils import get_llm_response
+from utils.epfl_meditron_utils import get_llm_response, gptq_model_options
 from utils.openai_utils import get_available_engines, get_search_query_type_options
 
 from sklearn.metrics import confusion_matrix, ConfusionMatrixDisplay
@@ -83,19 +83,18 @@ def display_streamlit_sidebar():
     st.sidebar.write('**Parameters**')
     form = st.sidebar.form("config_form", clear_on_submit=True)
 
-    model_option = form.selectbox("Quickly select a model", ("llama", "meditron"))
-    model_repo_id = form.text_input(label="Repo", value=model_option)#value=st.session_state["model_repo_id"])
-    model_filename = form.text_input(label="File name", value=st.session_state["model_filename"])
-    model_type = form.text_input(label="Model type", value=st.session_state["model_type"])
-    gpu_layers = form.slider('GPU Layers', min_value=0,
-                              max_value=100, value=st.session_state['gpu_layers'], step=1)
-
-    system_prompt = ""
-    #form.text_area(label='System prompt',
-    #                               value=st.session_state["system_prompt"])
+    model_name_or_path = form.selectbox("Select model", gptq_model_options())
+    
+    temperature = form.slider(label="Temperature", min_value=0.0, max_value=1.0, step=0.01, value=0.01)
+    do_sample = form.checkbox('do_sample')
+    top_p = form.slider(label="top_p", min_value=0.0, max_value=1.0, step=0.01, value=0.95)
+    top_k = form.slider(label="top_k", min_value=1, max_value=1000, step=1, value=40)
+    max_new_tokens = form.slider(label="max_new_tokens", min_value=32, max_value=512, step=1, value=32)
+    repetition_penalty = form.slider(label="repetition_penalty", min_value=0.0, max_value=1.0, step=0.01, value=0.95)
     
     temperature = form.slider('Temperature (0 = deterministic, 1 = more freedom)', min_value=0.0,
                               max_value=1.0, value=st.session_state['temperature'], step=0.1)
+    
     top_p = form.slider('top_p (0 = focused, 1 = broader answer range)', min_value=0.0,
                         max_value=1.0, value=st.session_state['top_p'], step=0.1)
 
@@ -104,22 +103,16 @@ def display_streamlit_sidebar():
     submitted = form.form_submit_button("Start session")
     if submitted and not st.session_state['session_started']:
         print('Parameters updated...')
-        restart_session()
         st.session_state['session_started'] = True
 
-        st.session_state["model_repo_id"] = model_repo_id
-        st.session_state["model_filename"] = model_filename
-        st.session_state["model_type"] = model_type
-        st.session_state['gpu_layers'] = gpu_layers
-
-        st.session_state["questions"] = []
-        st.session_state["lead_symptom"] = None
-        st.session_state["scenario_name"] = None
-        st.session_state["system_prompt"] = system_prompt
-        st.session_state['session_started'] = True
-        st.session_state["session_started"] = True
+        st.session_state["session_events"] = []
+        st.session_state["model_name_or_path"] = model_name_or_path
         st.session_state["temperature"] = temperature
+        st.session_state["do_sample"] = do_sample
         st.session_state["top_p"] = top_p
+        st.session_state["top_k"] = top_k
+        st.session_state["max_new_tokens"] = max_new_tokens
+        st.session_state["repetition_penalty"] = repetition_penalty
 
         st.rerun()
 
@@ -190,96 +183,20 @@ def get_chat_history_string(chat_history):
             raise Exception('Unknown role: ' + str(i["role"]))
 
     return res
-
-
-def restart_session():
-    print("Resetting params...")
-    st.session_state["emg_class_enabled"] = False
-    st.session_state["enable_llm_summary"] = False
-    st.session_state["num_variants"] = 3
-    st.session_state["lang_index"] = 0
-    st.session_state["llm_message"] = ""
-    st.session_state["llm_messages"] = []
-
-    st.session_state["triage_prompt_variants"] = ['''You are a telemedicine triage agent that decides between the following:
-    Emergency: Patient health is at risk if he doesn't speak to a Doctor urgently
-    Telecare: Patient can likely be treated remotely
-    General Practitioner: Patient should visit a GP for an ad-real consultation''',
-
-                                                  '''You are a Doctor assistant that decides if a medical case can likely be treated remotely by a Doctor or not.
-    The remote Doctor can write prescriptions and request the patient to provide a picture.
-    Provide the triage recommendation first, and then explain your reasoning respecting the format given below:
-    Treat remotely: <your reasoning>
-    Treat ad-real: <your reasoning>''',
-
-                                                  '''You are a medical triage agent working for the telemedicine Company Medgate based in Switzerland.
-    You decide if a case can be treated remotely or not, knowing that the remote Doctor can write prescriptions and request pictures.
-    Provide the triage recommendation first, and then explain your reasoning respecting the format given below:
-    Treat remotely: <your reasoning>
-    Treat ad-real: <your reasoning>''']
-
-    st.session_state['nbqs'] = []
-    st.session_state['citations'] = {}
-
-    st.session_state['past_messages'] = []
-    st.session_state["last_request"] = None
-    st.session_state["last_proposal"] = None
-
-    st.session_state['doctor_question'] = ''
-    st.session_state['patient_reply'] = ''
-
-    st.session_state['chat_history_array'] = []
-    st.session_state['chat_history'] = ''
-
-    st.session_state['feed_summary'] = ''
-    st.session_state['summary'] = ''
-
-    st.session_state["selected_guidelines"] = ["General"]
-    st.session_state["guidelines_dict"] = get_guidelines_dict()
-
-    st.session_state["triage_recommendation"] = ''
-
-    st.session_state["session_events"] = []
-
+      
 
 def init_session_state():
     print('init_session_state()')
     st.session_state['session_started'] = False
-    st.session_state['guidelines_ignored'] = False
-    st.session_state['model_index'] = 1
-
-    st.session_state["model_repo_id"] = "TheBloke/meditron-7B-GGUF"
-    st.session_state["model_filename"] = "meditron-7b.Q5_K_S.gguf"
-    st.session_state["model_type"] = "llama"
-    st.session_state['gpu_layers'] = 1
-
-    default_gender_index = 0
-    st.session_state['gender'] = get_genders()[default_gender_index]
-    st.session_state['gender_index'] = default_gender_index
-
-    st.session_state['age'] = 30
-
-    st.session_state['patient_medical_info'] = ''
-
-    default_search_query = 0
-    st.session_state['search_query_type'] = get_search_query_type_options()[default_search_query]
-    st.session_state['search_query_type_index'] = default_search_query
-    st.session_state['engine'] = get_available_engines()[0]
-    st.session_state['temperature'] = 0.0
-    st.session_state['top_p'] = 1.0
-    st.session_state['feed_chat_transcript'] = ''
-
-    st.session_state["llm_model"] = True
-    st.session_state["hugging_face_models"] = True
-    st.session_state["local_models"] = True
-
-    restart_session()
-
-    st.session_state['system_prompt'] = get_system_prompt()
-    st.session_state['system_prompt_after_on_change'] = get_system_prompt()
-
-    st.session_state["summary"] = ''
-
+    st.session_state["session_events"] = []
+    st.session_state["model_name_or_path"] = "TheBloke/meditron-7B-GPTQ"
+    st.session_state["temperature"] = 0.01
+    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["repetition_penalty"] = 1.1
+    st.session_state["system_message"] = "You are a medical expert that provides answers for a medically trained audience"  
 
 def get_genders():
     return ['Male', 'Female']
@@ -498,23 +415,55 @@ def get_diarized_f_path(audio_f_name):
     return DATA_FOLDER + base_name + ".txt"
 
 
+def get_prompt_format(model_name):
+    if model_name == "TheBloke/Llama-2-13B-chat-GPTQ":
+        return '''[INST] <<SYS>>
+        You are a helpful, respectful and honest assistant. Always answer as helpfully as possible, while being safe.  Your answers should not include any harmful, unethical, racist, sexist, toxic, dangerous, or illegal content. Please ensure that your responses are socially unbiased and positive in nature. If a question does not make any sense, or is not factually coherent, explain why instead of answering something not correct. If you don't know the answer to a question, please don't share false information.
+        <</SYS>>
+        {prompt}[/INST]
+
+        '''
+    if model_name == "TheBloke/Llama-2-7B-Chat-GPTQ":
+        return "[INST] <<SYS>>＼nYou are a helpful, respectful and honest assistant. Always answer as helpfully as possible, while being safe.  Your answers should not include any harmful, unethical, racist, sexist, toxic, dangerous, or illegal content. Please ensure that your responses are socially unbiased and positive in nature. If a question does not make any sense, or is not factually coherent, explain why instead of answering something not correct. If you don't know the answer to a question, please don't share false information.＼n<</SYS>>＼n{prompt}[/INST]"
+    
+    if model_name == "TheBloke/meditron-7B-GPTQ" or model_name == "TheBloke/meditron-70B-GPTQ":
+        return '''<|im_start|>system
+        {system_message}<|im_end|>
+        <|im_start|>user
+        {prompt}<|im_end|>
+        <|im_start|>assistant'''
+    
+    return ""
+
+def format_prompt(template, system_message, prompt):
+    if template == "":
+        return f"{system_message} {prompt}"
+    return template.format(system_message=system_message, prompt=prompt)
+
 def display_llm_output():
     st.header("LLM")
 
     form = st.form('llm')
 
-    llm_message = form.text_area('Message', value=st.session_state["llm_message"])
+    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"])
 
-    api_submitted = form.form_submit_button('Submit')
+    submitted = form.form_submit_button('Submit')
 
-    if api_submitted:
+    if submitted:
+        formatted_prompt = format_prompt(prompt_format, system_prompt, prompt)
+        print(f"Formatted prompt: {format_prompt}")
         llm_response = get_llm_response(
-            st.session_state["model_repo_id"],
-            st.session_state["model_filename"],
-            st.session_state["model_type"],
-            st.session_state["gpu_layers"],
-            "You are a medical assistant",
-            llm_message)
+            st.session_state["model_name"],
+            st.session_state["temperature"],
+            st.session_state["do_sample"],
+            st.session_state["top_p"],
+            st.session_state["top_k"],
+            st.session_state["max_new_tokens"],
+            st.session_state["repetition_penalty"],
+            formatted_prompt)
         st.write(llm_response)
         st.write('Done displaying LLM response')
 

utils/epfl_meditron_utils.py

@@ -1,49 +1,43 @@
 from transformers import AutoModelForCausalLM, AutoTokenizer, pipeline
 
-
-def get_llm_response(repo, filename, model_type, gpu_layers, system_message, prompt):
-
-    model_name_or_path = "TheBloke/meditron-7B-GPTQ"
-    # To use a different branch, change revision
-    # For example: revision="gptq-4bit-128g-actorder_True"
+def gptq_model_options():
+    return [
+        "TheBloke/Llama-2-7B-Chat-GPTQ", 
+        "TheBloke/Llama-2-13B-chat-GPTQ",
+        "TheBloke/meditron-7B-GPTQ",
+        "TheBloke/meditron-70B-GPTQ",
+    ]
+
+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,
                                                 device_map="auto",
                                                 trust_remote_code=False,
                                                 revision="main")
 
     tokenizer = AutoTokenizer.from_pretrained(model_name_or_path, use_fast=True)
-    
-    prompt_template=f'''<|im_start|>system
-        {system_message}<|im_end|>
-        <|im_start|>user
-        {prompt}<|im_end|>
-        <|im_start|>assistant
-        '''
-        
-    print("Template:")
-    print(prompt_template)
+            
+    print("Formatted prompt:")
+    print(formatted_prompt)
     
     print("\n\n*** Generate:")
 
-    input_ids = tokenizer(prompt_template, return_tensors='pt').input_ids.cuda()
-    output = model.generate(inputs=input_ids, temperature=0.01, do_sample=True, top_p=0.95, top_k=40, max_new_tokens=512)
-    print(tokenizer.decode(output[0])) 
+    input_ids = tokenizer(formatted_prompt, return_tensors='pt').input_ids.cuda()
+    output = model.generate(inputs=input_ids, temperature=temperature, do_sample=do_sample, top_p=top_p, top_k=top_k, max_new_tokens=max_new_tokens)
+    print(tokenizer.decode(output[0], skip_special_tokens=True)) 
 
     print("*** Pipeline:")
     pipe = pipeline(
         "text-generation",
         model=model,
         tokenizer=tokenizer,
-        max_new_tokens=512,
-        do_sample=True,
-        temperature=0.7,
-        top_p=0.95,
-        top_k=40,
-        repetition_penalty=1.1
+        max_new_tokens=max_new_tokens,
+        do_sample=do_sample,
+        temperature=temperature,
+        top_p=top_p,
+        top_k=top_p,
+        repetition_penalty=repetition_penalty
     )
 
-    
-
-    response = pipe(prompt_template)[0]['generated_text']
+    response = pipe(formatted_prompt)[0]['generated_text']
     print(response)
     return response