added gemma

app.py

@@ -9,6 +9,7 @@ import json
 
 retrieve_results = 10 
 show_examples = False
+llm_models_to_choose = ['mistralai/Mixtral-8x7B-Instruct-v0.1','mistralai/Mistral-7B-Instruct-v0.2', 'google/gemma-7b-it', 'None']
 
 generate_kwargs = dict(
     temperature = None,
@@ -56,11 +57,16 @@ def rag_cleaner(inp):
     date = inp['document_metadata']['_time']
     return f"{rank}. <b> {title} </b> \n Date : {date} \n Abstract: {content}"
 
-def get_prompt_text(question, context, formatted = True):
+def get_prompt_text(question, context, formatted = True, llm_model_picked = 'mistralai/Mistral-7B-Instruct-v0.2'):
     if formatted:
       sys_instruction = f"Context:\n {context} \n Given the following scientific paper abstracts, take a deep breath and lets think step by step to answer the question. Cite the titles of your sources when answering, do not cite links or dates."
       message = f"Question: {question}"
-      return f"<s>" + f"[INST] {sys_instruction} " +  f" {message} [/INST] </s> "
+      if 'mistralai' in llm_model_picked:
+          return f"<s>" + f"[INST] {sys_instruction}" +  f" {message}[/INST]"
+          
+      elif 'gemma' in llm_model_picked:
+          return f"<bos><start_of_turn>user\n{sys_instruction}" +  f" {message}<end_of_turn>\n"
+          
     return f"Context:\n {context} \n Given the following info, take a deep breath and lets think step by step to answer the question: {question}. Cite the titles of your sources when answering.\n\n"
 
 def get_references(question, retriever, k = retrieve_results):
@@ -76,7 +82,7 @@ with gr.Blocks(theme = gr.themes.Soft()) as demo:
       msg = gr.Textbox(label = 'Search', placeholder = 'What is Mistral?')
       with gr.Accordion("Advanced Settings", open=False):
         with gr.Row(equal_height = True):
-          llm_model = gr.Dropdown(choices = ['mistralai/Mixtral-8x7B-Instruct-v0.1','mistralai/Mistral-7B-Instruct-v0.2', 'None'], value = 'mistralai/Mistral-7B-Instruct-v0.2', label = 'LLM Model')
+          llm_model = gr.Dropdown(choices = llm_models_to_choose, value = 'mistralai/Mistral-7B-Instruct-v0.2', label = 'LLM Model')
           llm_results = gr.Slider(minimum=4, maximum=10, value=5, step=1, interactive=True, label="Top n results as context")
           stream_results = gr.Checkbox(value = True, label = "Stream output")
 
@@ -84,7 +90,7 @@ with gr.Blocks(theme = gr.themes.Soft()) as demo:
     input = gr.Textbox(show_label = False, visible = False)
     gr_md = gr.Markdown(mark_text + md_text_initial)
 
-    def update_with_rag_md(message, llm_results_use = 5):
+    def update_with_rag_md(message, llm_results_use = 5, llm_model_picked = 'mistralai/Mistral-7B-Instruct-v0.2'):
         rag_out = get_rag(message)
         md_text_updated = mark_text
         for i in range(retrieve_results):
@@ -99,7 +105,7 @@ with gr.Blocks(theme = gr.themes.Soft()) as demo:
           authors_formatted = f'*{authors}*' + ' \n\n'
         
           md_text_updated += paper_title + authors_formatted + paper_abs +  '\n---------------\n'+ '\n'
-        prompt = get_prompt_text(message, '\n\n'.join(rag_cleaner(out) for out in rag_out[:llm_results_use]))
+        prompt = get_prompt_text(message, '\n\n'.join(rag_cleaner(out) for out in rag_out[:llm_results_use]), llm_model_picked = llm_model_picked)
         return md_text_updated, prompt
 
     def ask_llm(prompt, llm_model_picked = 'mistralai/Mistral-7B-Instruct-v0.2', stream_outputs = False):
@@ -131,6 +137,6 @@ with gr.Blocks(theme = gr.themes.Soft()) as demo:
            return stream
 
 
-    msg.submit(update_with_rag_md, [msg, llm_results], [gr_md, input]).success(ask_llm, [input, llm_model, stream_results], output_text)
+    msg.submit(update_with_rag_md, [msg, llm_results, llm_model], [gr_md, input]).success(ask_llm, [input, llm_model, stream_results], output_text)
 
 demo.queue(default_concurrency_limit=10).launch()