New version using endpoints of Falcon and Mistral

app.py

@@ -2,10 +2,30 @@ import random
 import time
 import torch
 import gradio as gr
+import requests
 from transformers import AutoConfig, AutoTokenizer, AutoModelForSeq2SeqLM, AutoModelForCausalLM, MistralForCausalLM
 from peft import PeftModel, PeftConfig
 from textwrap import wrap, fill
 
+## using Falcon 7b Instruct
+Falcon_API_URL = "https://api-inference.huggingface.co/models/tiiuae/falcon-7b-instruct"
+HEADERS = {"Authorization": "Bearer <HF TOKEN>"}
+def falcon_query(payload):
+    response = requests.post(Falcon_API_URL, headers=HEADERS, json=payload)
+    return response.json()
+def falcon_inference(input_text):
+    payload = {"inputs": input_text}
+    return falcon_query(payload)
+
+## using Mistral
+Mistral_API_URL = "https://api-inference.huggingface.co/models/mistralai/Mistral-7B-v0.1"
+def mistral_query(payload):
+    response = requests.post(Mistral_API_URL , headers=HEADERS, json=payload)
+    return response.json()
+def mistral_inference(input_text):
+    payload = {"inputs": input_text}
+    return mistral_query(payload)
+
 # Functions to Wrap the Prompt Correctly
 def wrap_text(text, width=90):
     lines = text.split('\n')
@@ -65,8 +85,14 @@ class ChatbotInterface():
         
         self.submit.click(self.respond, [self.msg, self.chatbot], [self.msg, self.chatbot])
     
+    def respond(self, msg, chatbot):
+        raise NotImplementedError
+
+class GaiaMinimed(ChatbotInterface):
+    def __init__(self, name, system_prompt="You are an expert medical analyst that helps users with any medical related information."):
+        super().__init__(name, system_prompt)
+        
     def respond(self, msg, history):
-            #bot_message = random.choice(["Hello, I'm MedChat! How can I help you?", "Hello there! I'm Medchat, a medical assistant! How can I help you?"])
             formatted_input = f"{{{{ {self.system_prompt} }}}}\nUser: {msg}\n{self.name}:"
             input_ids = tokenizer.encode(
                 formatted_input, 
@@ -90,6 +116,26 @@ class ChatbotInterface():
 
             return "", self.chat_history
 
+class FalconBot(ChatbotInterface):
+    def __init__(self, name, system_prompt="You are an expert medical analyst that helps users with any medical related information."):
+        super().__init__(name, system_prompt)
+        
+    def respond(self, msg, chatbot):
+        falcon_response = falcon_inference(msg)
+        falcon_output = falcon_response[0]["generated_text"]
+        self.chat_history.append([msg, falcon_output])
+        return "", falcon_output
+
+class MistralBot(ChatbotInterface):
+    def __init__(self, name, system_prompt="You are an expert medical analyst that helps users with any medical related information."):
+        super().__init__(name, system_prompt)
+    
+    def respond(self, msg, chatbot):
+        mistral_response = mistral_inference(msg)
+        mistral_output = mistral_response[0]["generated_text"]
+        self.chat_history.append([msg, mistral_output])
+        return "", mistral_output
+
 if __name__ == "__main__":
     # Define the device
     device = "cuda" if torch.cuda.is_available() else "cpu"
@@ -120,15 +166,15 @@ if __name__ == "__main__":
         with gr.Row() as row:
             with gr.Column() as col1:
                 with gr.Tab("GaiaMinimed") as gaia:
-                    gaia_bot = ChatbotInterface("GaiaMinimed")
+                    gaia_bot = GaiaMinimed("GaiaMinimed")
             with gr.Column() as col2:
                 with gr.Tab("MistralMed") as mistral:
-                    mistral_bot = ChatbotInterface("MistralMed") 
+                    mistral_bot = MistralBot("MistralMed") 
                 with gr.Tab("Falcon-7B") as falcon7b:
-                    falcon_bot = ChatbotInterface("Falcon-7B")
+                    falcon_bot = FalconBot("Falcon-7B")
         
-        gaia_bot.msg.change(fn=lambda s: (s[::1], s[::1]), inputs=gaia_bot.msg, outputs=[mistral_bot.msg, falcon_bot.msg])
-        mistral_bot.msg.change(fn=lambda s: (s[::1], s[::1]), inputs=mistral_bot.msg, outputs=[gaia_bot.msg, falcon_bot.msg])
-        falcon_bot.msg.change(fn=lambda s: (s[::1], s[::1]), inputs=falcon_bot.msg, outputs=[gaia_bot.msg, mistral_bot.msg])
+        gaia_bot.msg.input(fn=lambda s: (s[::1], s[::1]), inputs=gaia_bot.msg, outputs=[mistral_bot.msg, falcon_bot.msg])
+        mistral_bot.msg.input(fn=lambda s: (s[::1], s[::1]), inputs=mistral_bot.msg, outputs=[gaia_bot.msg, falcon_bot.msg])
+        falcon_bot.msg.input(fn=lambda s: (s[::1], s[::1]), inputs=falcon_bot.msg, outputs=[gaia_bot.msg, mistral_bot.msg])
                 
     demo.launch()