update

pages/summarization_playground.py

@@ -12,6 +12,60 @@ import logging
 
 load_dotenv()
 
+custom_css = """
+gradio-app {
+    background: #eeeefc !important;
+}
+.bordered-text {
+    border-style: solid;
+    border-width: 1px;
+    padding: 5px;
+    margin-bottom: 0px;
+    border-radius: 1px;
+    font-family: Verdana;
+    font-size: 20px !important;
+    font-weight: bold ;
+    color:#000000;
+}
+.parameter-text {
+    border-style: solid;
+    border-width: 1px;
+    padding: 5px;
+    margin-bottom: 0px;
+    border-radius: 1px;
+    font-family: Verdana;
+    font-size: 10px !important;
+    font-weight: bold ;
+    color:#000000;
+}
+.title {
+    font-size: 35px;
+    font-weight: maroon;
+    font-family: Helvetica;
+}
+input-label {
+    font-size: 20px;
+    font-weight: bold;
+    font-family: Papyrus;
+}
+.custom-button {
+    background-color: white !important /* Green background */
+    color: black; /* White text */
+    border: none; /* Remove border */
+    padding: 10px 20px; /* Add padding */
+    text-align: center; /* Center text */
+    display: inline-block; /* Inline block */
+    font-size: 22px; /* Font size */
+    margin: 4px 2px; /* Margin */
+    cursor: pointer; /* Pointer cursor on hover */
+    border-radius: 4px; /* Rounded corners */
+}
+.custom-button:hover {
+    background-color: black;
+    color: white;
+}
+"""
+
 __model_on_gpu__ = ''
 model = {model_name: None for model_name in Model.__model_list__}
 
@@ -53,14 +107,14 @@ def get_model_batch_generation(model_name):
     return model[model_name]
 
 
-def generate_answer(sources, model_name, prompt):
+def generate_answer(sources, model_name, prompt, temperature, max_new_tokens, do_sample):
     model_device_check(model_name)
     content = prompt + '\n{' + sources + '}\n\nsummary:'
-    answer = model[model_name].gen(content)[0].strip()
+    answer = model[model_name].gen(content,temperature,max_new_tokens,do_sample)[0].strip()
 
     return answer
 
-def process_input(input_text, model_selection, prompt):
+def process_input(input_text, model_selection, prompt, temperature, max_new_tokens, do_sample):
     if input_text:
         logging.info("Start generation")
         response = generate_answer(input_text, model_selection, prompt)
@@ -75,13 +129,14 @@ def update_input(example):
     return examples[example]
 
 def create_summarization_interface():
-    with gr.Blocks(theme=gr.themes.Soft(spacing_size="sm",text_size="sm")) as demo:
+    with gr.Blocks(theme=gr.themes.Soft(spacing_size="sm",text_size="sm"), css=custom_css) as demo:
         gr.Markdown("## This is a playground to test prompts for clinical dialogue summarizations")
 
         with gr.Row():
             example_dropdown = gr.Dropdown(choices=list(examples.keys()), label="Choose an example", value=random_label)
             model_dropdown = gr.Dropdown(choices=Model.__model_list__, label="Choose a model", value=Model.__model_list__[0])
         
+        gr.Markdown("<div style='border: 4px solid white; padding: 3px; border-radius: 5px;width:100px;padding-top: 0.5px;padding-bottom: 10px;'><h3>Prompt 👥</h3></center></div>")
         Template_text = gr.Textbox(value="""Summarize the following dialogue""", label='Input Prompting Template', lines=8, placeholder='Input your prompts')
         datapoint = random.choice(dataset)
         input_text = gr.Textbox(label="Input Dialogue", lines=10, placeholder="Enter text here...", value=datapoint['section_text'] + '\n\nDialogue:\n' + datapoint['dialogue'])
@@ -89,24 +144,16 @@ def create_summarization_interface():
 
         with gr.Row():
             with gr.Row():
-                with gr.Column():
-                    gr.Markdown("<div style='border: 4px solid white; padding: 3px; border-radius: 5px;width:100px;padding-top: 0.5px;padding-bottom: 10px;'><h3>Prompt 👥</h3></center></div>")
-                    prompt = gr.Textbox(label="Input", lines=6, placeholder = "Enter the Patient-Doctor conversation here.",elem_classes=["bordered-text"])
-                    context = gr.Textbox(label="Context", placeholder="Enter relevant context about the patient medical history.",elem_classes="bordered-text")
-                    token = gr.Textbox(label="Token",elem_classes="bordered-text")
                 with gr.Column():
                     gr.Markdown("<div style='border: 4px solid white; padding: 2px; border-radius: 5px;width:130px;padding-bottom: 10px;'><b><h3>Parameters 📈</h3></center></b></div>")
-                    with gr.Column():
-                        topK = gr.Textbox(label="TopP",elem_classes="bordered-text")
-                        topP = gr.Textbox(label="TopK",elem_classes="bordered-text")
+                    with gr.Column(theme=gr.themes.Soft(spacing_size="sm",text_size="sm")):
                         temperature = gr.Textbox(label="Temperature",elem_classes="parameter-text")
                         max_new_tokens = gr.Textbox(label="Max New Tokens",elem_classes="parameter-text")
                         do_sample = gr.Dropdown(['Default','None'],label="Do Sample",elem_classes="parameter-text")
-                        return_text = gr.Dropdown(['Default','None'],label="Return Text",elem_classes="parameter-text")
             output = gr.Markdown(line_breaks=True)
 
         example_dropdown.change(update_input, inputs=[example_dropdown], outputs=[input_text])
-        submit_button.click(process_input, inputs=[input_text, model_dropdown, Template_text], outputs=[output])
+        submit_button.click(process_input, inputs=[input_text,model_dropdown,Template_text,temperature,max_new_tokens,do_sample], outputs=[output])
 
     return demo
 

utils/model.py

@@ -55,14 +55,14 @@ class Model(torch.nn.Module):
     def return_model(self):
         return self.model
     
-    def streaming(self, content_list, temp=0.001, max_length=500):
+    def streaming(self, content_list, temp=0.001, max_length=500, do_sample=True):
         # Convert list of texts to input IDs
         input_ids = self.tokenizer(content_list, return_tensors="pt", padding=True, truncation=True).input_ids.to(self.model.device)
 
         # Set up the initial generation parameters
         gen_kwargs = {
             "input_ids": input_ids,
-            "do_sample": True,
+            "do_sample": do_sample,
             "temperature": temp,
             "eos_token_id": self.tokenizer.eos_token_id,
             "max_new_tokens": 1,  # Generate one token at a time
@@ -96,7 +96,7 @@ class Model(torch.nn.Module):
                 gen_kwargs["input_ids"] = gen_kwargs["input_ids"][active_sequences]
 
 
-    def gen(self, content_list, temp=0.001, max_length=500):
+    def gen(self, content_list, temp=0.001, max_length=500, do_sample=True):
         # Convert list of texts to input IDs
         input_ids = self.tokenizer(content_list, return_tensors="pt", padding=True, truncation=True).input_ids.to(self.model.device)
 
@@ -104,7 +104,7 @@ class Model(torch.nn.Module):
         outputs = self.model.generate(
             input_ids,
             max_new_tokens=max_length,
-            do_sample=True,
+            do_sample=do_sample,
             temperature=temp,
             eos_token_id=self.tokenizer.eos_token_id,
         )