Refining UI

app.py

@@ -4,6 +4,7 @@ import google.generativeai as genai
 import openai
 from collections import Counter
 from huggingface_hub import InferenceClient
+import re
 
 def api_check_msg(api_key, selected_model):
     res = validate_api_key(api_key, selected_model)
@@ -96,6 +97,21 @@ def generate_text_llama(key, prompt, temperature, top_p):
         return response[len(prompt):]
     return response
 
+def sanitize_outputs(outputs):
+    sanitized_results = []
+    for output in outputs:
+        output = output.replace("\n", " ")
+        output = re.sub(r"(Diagnose:|Answer:)", "", output, flags=re.IGNORECASE).strip()
+        diagnoses = ["Psoriasis", "Arthritis", "Bronchial asthma", "Cervical spondylosis"]
+        found_diagnoses = [disease for disease in diagnoses if disease in output]
+
+        if found_diagnoses:
+            sanitized_results.append(found_diagnoses[0])
+        else:
+            sanitized_results.append("Unknown")  # Handle case where no valid diagnosis is found
+
+    return sanitized_results
+
 
 def diagnose(gpt_key, llama_key, gemini_key, top_p, temperature, symptoms):
 
@@ -110,6 +126,7 @@ def diagnose(gpt_key, llama_key, gemini_key, top_p, temperature, symptoms):
         gemini_message = generate_text_gemini(gemini_key, prompt, temperature, top_p)
 
         outputs = [gpt_message, llama_message, gemini_message]
+        outputs = sanitize_outputs(outputs)
         output_counts = Counter(outputs)
         majority_output, majority_count = output_counts.most_common(1)[0]
         confidence = (majority_count / len(outputs)) * 100