Update app.py

app.py

@@ -7,26 +7,25 @@ from transformers import AutoConfig, AutoTokenizer, AutoModelForSeq2SeqLM, AutoM
 from peft import PeftModel, PeftConfig
 import torch
 import gradio as gr
-# Functions to Wrap the Prompt Correctly
 
+# Functions to Wrap the Prompt Correctly
 def wrap_text(text, width=90):
     lines = text.split('\n')
     wrapped_lines = [textwrap.fill(line, width=width) for line in lines]
     wrapped_text = '\n'.join(wrapped_lines)
     return wrapped_text
 
-def multimodal_prompt(input_text, system_prompt="", max_length=512):
+def multimodal_prompt(user_input, system_prompt="You are an expert medical analyst:"):
     """
-    Generates text using a large language model, given a prompt and a device.
+    Generates text using a large language model, given a user input and a system prompt.
     Args:
-        input_text: The input text to generate a response for.
+        user_input: The user's input text to generate a response for.
         system_prompt: Optional system prompt.
-        max_length: Maximum length of the generated text.
     Returns:
         A string containing the generated text.
     """
-    # Modify the input text to include the desired format
-    formatted_input = f"""<s>[INST]{input_text}[/INST]"""
+    # Combine user input and system prompt
+    formatted_input = f"<s>[INST]{system_prompt} {user_input}[/INST]"
 
     # Encode the input text
     encodeds = tokenizer(formatted_input, return_tensors="pt", add_special_tokens=False)
@@ -50,7 +49,6 @@ def multimodal_prompt(input_text, system_prompt="", max_length=512):
 
     return response_text
 
-
 # Define the device
 device = "cuda" if torch.cuda.is_available() else "cpu"
 
@@ -64,7 +62,6 @@ tokenizer = AutoTokenizer.from_pretrained("mistralai/Mistral-7B-v0.1", trust_rem
 tokenizer.pad_token = tokenizer.eos_token
 tokenizer.padding_side = 'left'
 
-
 # Specify the configuration class for the model
 #model_config = AutoConfig.from_pretrained(base_model_id)
 
@@ -80,9 +77,12 @@ class ChatBot:
     def __init__(self):
         self.history = []
 
-    def predict(self, input_text):
+    def predict(self, user_input, system_prompt="You are an expert medical analyst:"):
+        # Combine user input and system prompt
+        formatted_input = f"<s>[INST]{system_prompt} {user_input}[/INST]"
+
         # Encode user input
-        user_input_ids = tokenizer.encode(input_text, return_tensors="pt")
+        user_input_ids = tokenizer.encode(formatted_input, return_tensors="pt")
 
         # Concatenate the user input with chat history
         if len(self.history) > 0:
@@ -104,15 +104,14 @@ bot = ChatBot()
 
 title = "👋🏻Welcome to Tonic's MistralMed Chat🚀"
 description = "You can use this Space to test out the current model (MistralMed) or duplicate this Space and use it for any other model on 🤗HuggingFace. Join me on Discord to build together."
-examples = [["<s>[INST] What is the proper treatment for bucal herpes?[/INST]"]]
-
+examples = [["What is the proper treatment for buccal herpes?"]]
 
 iface = gr.Interface(
     fn=bot.predict,
     title=title,
     description=description,
     examples=examples,
-    inputs="text",
+    inputs=["text", "text"],  # Take user input and system prompt separately
     outputs="text",
     theme="ParityError/Anime"
 )