Added prompt templates to output correct diagnoses (#2)

- Added prompt templates to output correct diagnoses (373220b9e98f8db692fd95aa4aac60e91d806282)


Co-authored-by: Bharat Raghunathan <bharat-raghunathan@users.noreply.huggingface.co>

app.py

@@ -1,171 +1,184 @@
-import gradio as gr
-import os
-import requests
-from huggingface_hub import InferenceClient
-import google.generativeai as genai
-import openai
-
-def api_check_msg(api_key, selected_model):
-    res = validate_api_key(api_key, selected_model)
-    return res["message"]
-
-def validate_api_key(api_key, selected_model):
-    # Check if the API key is valid for GPT-3.5-Turbo
-    if "GPT" in selected_model:
-        url = "https://api.openai.com/v1/models"
-        headers = {
-            "Authorization": f"Bearer {api_key}"
-        }
-        try:
-            response = requests.get(url, headers=headers)
-            if response.status_code == 200:
-                return {"is_valid": True, "message": '<p style="color: green;">API Key is valid!</p>'}
-            else:
-                return {"is_valid": False, "message": f'<p style="color: red;">Invalid OpenAI API Key. Status code: {response.status_code}</p>'}
-        except requests.exceptions.RequestException as e:
-            return {"is_valid": False, "message": f'<p style="color: red;">Invalid OpenAI API Key. Error: {e}</p>'}
-    elif "Llama" in selected_model:
-        url = "https://huggingface.co/api/whoami-v2"
-        headers = {
-            "Authorization": f"Bearer {api_key}"
-        }
-        try:
-            response = requests.get(url, headers=headers)
-            if response.status_code == 200:
-                return {"is_valid": True, "message": '<p style="color: green;">API Key is valid!</p>'}
-            else:
-                return {"is_valid": False, "message": f'<p style="color: red;">Invalid Hugging Face API Key. Status code: {response.status_code}</p>'}
-        except requests.exceptions.RequestException as e:
-            return {"is_valid": False, "message": f'<p style="color: red;">Invalid Hugging Face API Key. Error: {e}</p>'}
-    elif "Gemini" in selected_model:
-        try:
-            genai.configure(api_key=api_key)
-            model = genai.GenerativeModel("gemini-1.5-flash")
-            response = model.generate_content("Help me diagnose the patient.")
-            return {"is_valid": True, "message": '<p style="color: green;">API Key is valid!</p>'}
-        except Exception as e:
-            return {"is_valid": False, "message": f'<p style="color: red;">Invalid Google API Key. Error: {e}</p>'}
-
-def generate_text_chatgpt(key, prompt, temperature, top_p):
-
-    openai.api_key = key
-
-    response = openai.chat.completions.create(
-      model="gpt-3.5-turbo-1106",
-      messages=[{"role": "system", "content": "You are a talented diagnostician who is diagnosing a patient."},
-                {"role": "user", "content": prompt}],
-      temperature=temperature,
-      max_tokens=50,
-      top_p=top_p,
-      frequency_penalty=0
-    )
-
-    return response.choices[0].message.content
-
-
-def generate_text_gemini(key, prompt, temperature, top_p):
-    genai.configure(api_key=key)
-
-    generation_config = genai.GenerationConfig(
-        max_output_tokens=len(prompt)+50,
-        temperature=temperature,
-        top_p=top_p,
-    )
-    model = genai.GenerativeModel("gemini-1.5-flash", generation_config=generation_config)
-    response = model.generate_content(prompt)
-    return response.text
-
-
-def generate_text_llama(key, prompt, temperature, top_p):
-    model_name = "meta-llama/Meta-Llama-3-8B-Instruct"
-    client = InferenceClient(api_key=key)
-
-    messages = [{"role": "system", "content": "You are a talented diagnostician who is diagnosing a patient."},
-                {"role": "user","content": prompt}]
-
-    completion = client.chat.completions.create(
-        model=model_name,
-        messages=messages, 
-        max_tokens=len(prompt)+50,
-        temperature=temperature,
-        top_p=top_p
-    )
-
-    response = completion.choices[0].message.content
-    if len(response) > len(prompt):
-        return response[len(prompt):]
-    return response
-
-
-def diagnose(key, model, top_k, temperature, symptom_prompt):
-
-    model_map = {
-        "GPT-3.5-Turbo": "GPT",
-        "Llama-3": "Llama",
-        "Gemini-1.5": "Gemini"
-    }
-    if symptom_prompt:
-        if "GPT" in model:
-            message = generate_text_chatgpt(key, symptom_prompt, temperature, top_k)
-        elif "Llama" in model:
-            message = generate_text_llama(key, symptom_prompt, temperature, top_k)
-        elif "Gemini" in model:
-            message = generate_text_gemini(key, symptom_prompt, temperature, top_k)
-        else:
-            message = "Incorrect model, please try again."
-    else:
-        message = "Please add the symptoms data"
-    
-    return message
-
-def update_model_components(selected_model):
-    model_map = {
-                "GPT-3.5-Turbo": "GPT",
-                "Llama-3": "Llama",
-                "Gemini-1.5": "Gemini"
-            }
-
-    link_map = {
-        "GPT-3.5-Turbo": "https://platform.openai.com/account/api-keys",
-        "Llama-3": "https://hf.co/settings/tokens",
-        "Gemini-1.5": "https://aistudio.google.com/apikey"
-    }
-    textbox_label = f"Please input the API key for your {model_map[selected_model]} model"
-    button_value = f"Don't have an API key? Get one for the {model_map[selected_model]} model here."
-    button_link = link_map[selected_model]
-    return gr.update(label=textbox_label), gr.update(value=button_value, link=button_link)
-
-def toggle_button(symptoms_text, api_key, model):
-    if symptoms_text.strip() and validate_api_key(api_key, model):
-        return gr.update(interactive=True)
-    return gr.update(interactive=False)
-
-
-with gr.Blocks() as ui:
-
-    with gr.Row(equal_height=500):
-        with gr.Column(scale=1, min_width=300):
-            model = gr.Radio(label="LLM Selection", value="GPT-3.5-Turbo",  
-                             choices=["GPT-3.5-Turbo", "Llama-3", "Gemini-1.5"])
-            is_valid = False
-            key = gr.Textbox(label="Please input the API key for your Large Language model", type="password")
-            status_message = gr.HTML(label="Validation Status")
-            key.input(fn=api_check_msg, inputs=[key, model], outputs=status_message)
-            button = gr.Button(value="Don't have an API key? Get one for the GPT model here.", link="https://platform.openai.com/account/api-keys")
-            model.change(update_model_components, inputs=model, outputs=[key, button])
-            gr.ClearButton(key, variant="primary")
-            
-        with gr.Column(scale=2, min_width=600):
-            gr.Markdown("## Hello, Welcome to the GUI by Team #9.")
-            temperature = gr.Slider(0.0, 1.0, value=0.7, step = 0.05, label="Temperature", info="Set the Temperature")
-            top_p = gr.Slider(0.0, 1.0, value=0.9, step = 0.05, label="top-p value", info="Set the sampling nucleus parameter")
-            symptoms = gr.Textbox(label="Add the symptom data in the input to receive diagnosis")
-            llm_btn = gr.Button(value="Diagnose Disease", variant="primary", elem_id="diagnose", interactive=False)
-            symptoms.input(toggle_button, inputs=[symptoms, key, model], outputs=llm_btn)
-            key.input(toggle_button, inputs=[symptoms, key, model], outputs=llm_btn)
-            model.change(toggle_button, inputs=[symptoms, key, model], outputs=llm_btn)
-            output = gr.Textbox(label="LLM Output Status", interactive=False, placeholder="Output will appear here...")
-            llm_btn.click(fn=diagnose, inputs=[key, model, top_p, temperature, symptoms], outputs=output, api_name="LLM_Comparator")
-
-
+import gradio as gr
+import os
+import requests
+from huggingface_hub import InferenceClient
+import google.generativeai as genai
+import openai
+
+def api_check_msg(api_key, selected_model):
+    res = validate_api_key(api_key, selected_model)
+    return res["message"]
+
+def validate_api_key(api_key, selected_model):
+    # Check if the API key is valid for GPT-3.5-Turbo
+    if "GPT" in selected_model:
+        url = "https://api.openai.com/v1/models"
+        headers = {
+            "Authorization": f"Bearer {api_key}"
+        }
+        try:
+            response = requests.get(url, headers=headers)
+            if response.status_code == 200:
+                return {"is_valid": True, "message": '<p style="color: green;">API Key is valid!</p>'}
+            else:
+                return {"is_valid": False, "message": f'<p style="color: red;">Invalid OpenAI API Key. Status code: {response.status_code}</p>'}
+        except requests.exceptions.RequestException as e:
+            return {"is_valid": False, "message": f'<p style="color: red;">Invalid OpenAI API Key. Error: {e}</p>'}
+    elif "Llama" in selected_model:
+        url = "https://huggingface.co/api/whoami-v2"
+        headers = {
+            "Authorization": f"Bearer {api_key}"
+        }
+        try:
+            response = requests.get(url, headers=headers)
+            if response.status_code == 200:
+                return {"is_valid": True, "message": '<p style="color: green;">API Key is valid!</p>'}
+            else:
+                return {"is_valid": False, "message": f'<p style="color: red;">Invalid Hugging Face API Key. Status code: {response.status_code}</p>'}
+        except requests.exceptions.RequestException as e:
+            return {"is_valid": False, "message": f'<p style="color: red;">Invalid Hugging Face API Key. Error: {e}</p>'}
+    elif "Gemini" in selected_model:
+        try:
+            genai.configure(api_key=api_key)
+            model = genai.GenerativeModel("gemini-1.5-flash")
+            response = model.generate_content("Help me diagnose the patient.")
+            return {"is_valid": True, "message": '<p style="color: green;">API Key is valid!</p>'}
+        except Exception as e:
+            return {"is_valid": False, "message": f'<p style="color: red;">Invalid Google API Key. Error: {e}</p>'}
+
+def generate_text_chatgpt(key, prompt, temperature, top_p):
+
+    openai.api_key = key
+
+    prompt_template = f"""
+    {prompt} <Choose only one among the words Psoriasis, Arthritis, Bronchial asthma or Cervical spondylosis>
+    """
+
+    response = openai.chat.completions.create(
+      model="gpt-3.5-turbo-1106",
+      messages=[{"role": "system", "content": "You are a talented diagnostician who is diagnosing a patient."},
+                {"role": "user", "content": prompt_template}],
+      temperature=temperature,
+      max_tokens=50,
+      top_p=top_p,
+      frequency_penalty=0
+    )
+
+    return response.choices[0].message.content
+
+
+def generate_text_gemini(key, prompt, temperature, top_p):
+    genai.configure(api_key=key)
+
+    prompt_template = f"""
+    {prompt} <Choose only one among the words Psoriasis, Arthritis, Bronchial asthma or Cervical spondylosis>
+    """
+
+    generation_config = genai.GenerationConfig(
+        max_output_tokens=len(prompt_template)+50,
+        temperature=temperature,
+        top_p=top_p,
+    )
+    model = genai.GenerativeModel("gemini-1.5-flash", generation_config=generation_config)
+    response = model.generate_content(prompt_template)
+    return response.text
+
+
+def generate_text_llama(key, prompt, temperature, top_p):
+    model_name = "meta-llama/Meta-Llama-3-8B-Instruct"
+    client = InferenceClient(api_key=key)
+
+    prompt_template = f"""
+    {prompt} <Choose only one among the words Psoriasis, Arthritis, Bronchial asthma or Cervical spondylosis>
+    Do not list the symptoms again in the response. Do not add any additional text. Do not attempt to explain your answer.
+    """
+
+    messages = [{"role": "system", "content": "You are a talented diagnostician who is diagnosing a patient."},
+                {"role": "user","content": prompt_template}]
+
+    completion = client.chat.completions.create(
+        model=model_name,
+        messages=messages, 
+        max_tokens=len(prompt_template)+50,
+        temperature=temperature,
+        top_p=top_p
+    )
+
+    response = completion.choices[0].message.content
+    if len(response) > len(prompt_template):
+        return response[len(prompt_template):]
+    return response
+
+
+def diagnose(key, model, top_k, temperature, symptom_prompt):
+
+    model_map = {
+        "GPT-3.5-Turbo": "GPT",
+        "Llama-3": "Llama",
+        "Gemini-1.5": "Gemini"
+    }
+    if symptom_prompt:
+        if "GPT" in model:
+            message = generate_text_chatgpt(key, symptom_prompt, temperature, top_k)
+        elif "Llama" in model:
+            message = generate_text_llama(key, symptom_prompt, temperature, top_k)
+        elif "Gemini" in model:
+            message = generate_text_gemini(key, symptom_prompt, temperature, top_k)
+        else:
+            message = "Incorrect model, please try again."
+    else:
+        message = "Please add the symptoms data"
+    
+    return message
+
+def update_model_components(selected_model):
+    model_map = {
+                "GPT-3.5-Turbo": "GPT",
+                "Llama-3": "Llama",
+                "Gemini-1.5": "Gemini"
+            }
+
+    link_map = {
+        "GPT-3.5-Turbo": "https://platform.openai.com/account/api-keys",
+        "Llama-3": "https://hf.co/settings/tokens",
+        "Gemini-1.5": "https://aistudio.google.com/apikey"
+    }
+    textbox_label = f"Please input the API key for your {model_map[selected_model]} model"
+    button_value = f"Don't have an API key? Get one for the {model_map[selected_model]} model here."
+    button_link = link_map[selected_model]
+    return gr.update(label=textbox_label), gr.update(value=button_value, link=button_link)
+
+def toggle_button(symptoms_text, api_key, model):
+    if symptoms_text.strip() and validate_api_key(api_key, model):
+        return gr.update(interactive=True)
+    return gr.update(interactive=False)
+
+
+with gr.Blocks() as ui:
+
+    with gr.Row(equal_height=500):
+        with gr.Column(scale=1, min_width=300):
+            model = gr.Radio(label="LLM Selection", value="GPT-3.5-Turbo",  
+                             choices=["GPT-3.5-Turbo", "Llama-3", "Gemini-1.5"])
+            is_valid = False
+            key = gr.Textbox(label="Please input the API key for your Large Language model", type="password")
+            status_message = gr.HTML(label="Validation Status")
+            key.input(fn=api_check_msg, inputs=[key, model], outputs=status_message)
+            button = gr.Button(value="Don't have an API key? Get one for the GPT model here.", link="https://platform.openai.com/account/api-keys")
+            model.change(update_model_components, inputs=model, outputs=[key, button])
+            gr.ClearButton(key, variant="primary")
+            
+        with gr.Column(scale=2, min_width=600):
+            gr.Markdown("## Hello, Welcome to the GUI by Team #9.")
+            temperature = gr.Slider(0.0, 1.0, value=0.7, step = 0.05, label="Temperature", info="Set the Temperature")
+            top_p = gr.Slider(0.0, 1.0, value=0.9, step = 0.05, label="top-p value", info="Set the sampling nucleus parameter")
+            symptoms = gr.Textbox(label="Add the symptom data in the input to receive diagnosis")
+            llm_btn = gr.Button(value="Diagnose Disease", variant="primary", elem_id="diagnose", interactive=False)
+            symptoms.input(toggle_button, inputs=[symptoms, key, model], outputs=llm_btn)
+            key.input(toggle_button, inputs=[symptoms, key, model], outputs=llm_btn)
+            model.change(toggle_button, inputs=[symptoms, key, model], outputs=llm_btn)
+            output = gr.Textbox(label="LLM Output Status", interactive=False, placeholder="Output will appear here...")
+            llm_btn.click(fn=diagnose, inputs=[key, model, top_p, temperature, symptoms], outputs=output, api_name="LLM_Comparator")
+
+
 ui.launch(share=True)