Update app.py

app.py

@@ -3,12 +3,18 @@ import torch
 from transformers import AutoModelForCausalLM, AutoTokenizer
 import spaces
 
-# Initialize model
+# Initialize device
 device = "cuda" if torch.cuda.is_available() else "cpu"
-model_name = "linjc16/Panacea-7B-Chat"
-model = AutoModelForCausalLM.from_pretrained(model_name, torch_dtype=torch.float16)
+
+# Load model names from an external file
+with open("models.txt", "r") as f:
+    model_list = [line.strip() for line in f.readlines()]
+
+# Load default model
+current_model_name = model_list[0]
+model = AutoModelForCausalLM.from_pretrained(current_model_name, torch_dtype=torch.float16)
 model.to(device)
-tokenizer = AutoTokenizer.from_pretrained(model_name)
+tokenizer = AutoTokenizer.from_pretrained(current_model_name)
 
 # Default values for system and user input
 test_instruction_string = """
@@ -53,6 +59,12 @@ PRIOR CONCURRENT THERAPY:
 * No prior radiotherapy to \> 30% of the bone marrow or more than standard adjuvant pelvic radiotherapy for rectal cancer <Conditions:>Lung Cancer, Unspecified Adult Solid Tumor, Protocol Specific,  <Interventions:>indocyanine green, lidocaine, vinorelbine ditartrate, high performance liquid chromatography, intracellular fluorescence polarization analysis, liquid chromatography, mass spectrometry, pharmacological study <StudyType:>INTERVENTIONAL <PrimaryOutcomes:>Area Under the Curve, Number of Participants With Grade 3 and 4 Toxicities <OverallStatus:>COMPLETED
 """
 
+def load_model(model_name):
+    global model, tokenizer
+    model = AutoModelForCausalLM.from_pretrained(model_name, torch_dtype=torch.float16)
+    model.to(device)
+    tokenizer = AutoTokenizer.from_pretrained(model_name)
+
 @spaces.GPU
 def generate_response(system_instruction, user_input):
     # Format the prompt using the messages structure
@@ -67,15 +79,18 @@ def generate_response(system_instruction, user_input):
     with torch.no_grad():
         generated_ids = model.generate(model_inputs, max_new_tokens=1000, do_sample=True)
     # Extract only the bot's response, omitting the prompt
-    response = tokenizer.batch_decode(generated_ids, skip_special_tokens=True)[0].split("\n")[-2:]
+    response = tokenizer.batch_decode(generated_ids, skip_special_tokens=True)[0].split("\n")[-1].strip()
 
     return response
 
 # Gradio interface setup
 with gr.Blocks() as demo:
-    gr.Markdown("# Clinical Trial Chatbot")
+    gr.Markdown("# Clinical Trial Chatbot with Model Selection")
 
     with gr.Row():
+        # Dropdown for selecting model
+        model_dropdown = gr.Dropdown(choices=model_list, value=current_model_name, label="Select Model")
+        
         # Left sidebar for inputs
         with gr.Column():
             system_instruction = gr.Textbox(
@@ -96,7 +111,8 @@ with gr.Blocks() as demo:
                 label="Bot Response", interactive=False, placeholder="Response will appear here."
             )
 
-    # Link submit button to the generate_response function
+    # Link model selection and submit button to functions
+    model_dropdown.change(lambda m: load_model(m), inputs=model_dropdown, outputs=[])
     submit_btn.click(generate_response, [system_instruction, user_input], response_display)
 
 # Launch the app