Update app.py

app.py

@@ -13,37 +13,26 @@ tokenizer = AutoTokenizer.from_pretrained(model_name)
 # 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 = """
 ###Question:
-
 <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³
@@ -57,9 +46,7 @@ PATIENT CHARACTERISTICS:
 * 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)
@@ -79,7 +66,8 @@ def generate_response(system_instruction, user_input):
     # Generate model response
     with torch.no_grad():
         generated_ids = model.generate(model_inputs, max_new_tokens=1000, do_sample=True)
-    response = tokenizer.batch_decode(generated_ids, skip_special_tokens=True)[0].split("Bot:")[-1].strip()
+    # Extract only the bot's response, omitting the prompt
+    response = tokenizer.batch_decode(generated_ids, skip_special_tokens=True)[0].split("\n")[-1].strip()
 
     return response