Update app.py

app.py

@@ -10,90 +10,54 @@ device = "cuda" if torch.cuda.is_available() else "cpu"
 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(current_model_name)
+# Define variables for the model and tokenizer
+model = None
+tokenizer = None
 
-# Default values for system and user input
-test_instruction_string = """
-You are a helpful assistant with experience in the clinical domain and clinical trial design. You'll be asked queries related to clinical trials. These inquiries will be delineated by a '##Question' heading. Inside these queries, expect to find comprehensive details about the clinical trial structured within specific subsections, indicated by '<>' tags. These subsections include essential information such as the trial's title, brief summary, condition under study, inclusion and exclusion criteria, intervention, and outcomes.
-In answer to this question, return a list of probable baseline features (each feature should be enclosed within a pair of backticks and each feature should be separated by commas from other features) of the clinical trial. Baseline features are the set of baseline or demographic characteristics that are assessed at baseline and used in the analysis of the primary outcome measure(s) to characterize the study population and assess validity. Clinical trial-related publications typically include a table of baseline features assessed  by arm or comparison group and for the entire population of participants in the clinical trial.
-Do not give any additional explanations or use any tags or headers, only return the list of baseline features.
-"""
-
-test_input_string = """
-<Title:>Vinorelbine in Treating Patients With Advanced Solid Tumors That Have Not Responded to Treatment and Liver Dysfunction <BriefSummary:>RATIONALE: Drugs used in chemotherapy, such as vinorelbine, work in different ways to stop the growth of tumor cells, either by killing the cells or by stopping them from dividing.
-PURPOSE: This pilot trial is studying the side effects and best dose of vinorelbine in treating patients with advanced solid tumors that have not responded to treatment and liver dysfunction. <EligibilityCriteria:>DISEASE CHARACTERISTICS:
-* Histologically confirmed advanced solid tumor
-  * Any histology allowed
-* Refractory to standard therapy OR no standard therapy exists
-  * Previously untreated non-small cell lung cancer allowed, provided abnormal liver function is present, defined as moderate (group 3) or severe (group 4)
-* Measurable disease not required
-  * Present measurable disease requires baseline measurements within 4 weeks of study entry
-* Patients with acute hepatitis from viral or drug etiologies should recover to a stable baseline prior to study therapy
-* History of brain metastasis allowed, provided the following criteria are met:
-  * Metastasis has been controlled by radiotherapy or surgery
-  * Patient is not currently on corticosteroids
-  * Neurologic status is stable
-PATIENT CHARACTERISTICS:
-* Karnofsky performance status 70-100%
-* Life expectancy ≥ 2 months
-* ANC = 1,500/mm³
-* Platelet count = 100,000/mm³
-* Hemoglobin = 10 g/dL (transfusion to this level allowed)
-* Creatinine \< 1.5 mg/dL OR creatinine clearance \> 60 mL/ min
-* Patients with EKG evidence of first- or second-degree AV block or left or right bundle branch block are ineligible for the lidocaine bolus, but may otherwise be treated on this protocol
-* Not pregnant or nursing
-* Negative pregnancy test
-* Fertile patients must use effective contraception
-* No concurrent illness (e.g., cardiovascular, pulmonary, or central nervous system) that is poorly controlled or of such severity that the investigator deems unwise to enter the patient on protocol
-* Must have ability to comply with study treatment and required tests
-* Obstructive jaundice requires a drainage procedure prior to study treatment
-PRIOR CONCURRENT THERAPY:
-* See Disease Characteristics
-* Recovered from prior therapy
-* At least 3 weeks since prior chemotherapy (6 weeks for mitomycin or nitrosourea 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
-"""
-
-# Adjust load_model to ensure models are loaded to the correct device on demand
 def load_model(model_name):
     global model, tokenizer
     device = "cuda" if torch.cuda.is_available() else "cpu"
     
-    # Remove previous model from GPU memory if switching models
-    if 'model' in globals():
+    # Clear previous model from memory
+    if model:
         del model
         torch.cuda.empty_cache()
     
-    model = AutoModelForCausalLM.from_pretrained(model_name, torch_dtype=torch.float16).to(device)
-    tokenizer = AutoTokenizer.from_pretrained(model_name)
+    try:
+        # Load the specified model and move it to the device
+        model = AutoModelForCausalLM.from_pretrained(model_name, torch_dtype=torch.float16).to(device)
+        tokenizer = AutoTokenizer.from_pretrained(model_name)
+        print(f"Loaded model: {model_name}")
+    except ValueError as e:
+        print(f"Failed to load model '{model_name}': {e}")
+        model = None  # Reset if loading fails
+        tokenizer = None
+
+# Load the initial model
+current_model_name = model_list[0]
+load_model(current_model_name)
 
+# Default values for system and user input
+test_instruction_string = """You are a helpful assistant with experience in the clinical domain..."""
+test_input_string = """<Title:>Vinorelbine in Treating Patients With Advanced Solid Tumors..."""
 
 alpaca_prompt = """Below is an instruction that describes a task, paired with an input that provides further context. Write a response that appropriately completes the request.
-
 ### Instruction:
 {}
-
 ### Input:
 {}
-
 ### Response:
 {}"""
 
 @spaces.GPU
 def generate_response(system_instruction, user_input):
-    # Determine the correct device
-    device = "cuda" if torch.cuda.is_available() else "cpu"
-    model.to(device)
+    # Check if model is loaded
+    if model is None:
+        return "Model is not loaded. Please select a compatible model."
 
-    # Prepare the input in the appropriate format and move it to the correct device
-    inputs = tokenizer(
-        [alpaca_prompt.format(system_instruction, user_input, "")], 
-        return_tensors="pt"
-    ).to(device)
+    # Prepare input and move to the appropriate device
+    device = model.device
+    inputs = tokenizer([alpaca_prompt.format(system_instruction, user_input, "")], return_tensors="pt").to(device)
 
     # Define generation configuration
     meta_config = {
@@ -120,7 +84,6 @@ with gr.Blocks() as demo:
     gr.Markdown("# Clinical Trial Chatbot with Model Selection")
 
     with gr.Row():
-        # Left column for user inputs
         with gr.Column():
             system_instruction = gr.Textbox(
                 value=test_instruction_string,
@@ -134,7 +97,6 @@ with gr.Blocks() as demo:
             )
             submit_btn = gr.Button("Submit")
 
-        # Right column for model selection and bot response
         with gr.Column():
             model_dropdown = gr.Dropdown(choices=model_list, value=current_model_name, label="Select Model")
             response_display = gr.Textbox(
@@ -145,5 +107,5 @@ with gr.Blocks() as demo:
     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
-demo.launch()
+# Launch the app with share=True to make it public
+demo.launch(share=True)