Address feedback

Signed-off-by: SimJeg <[email protected]>

app.py

@@ -30,13 +30,13 @@ def process_request(url, question, press_name, compression_ratio):
     """ """
 
     if press_name not in press_dict:
-        return f"Invalid press type selected: {press_name}", -1
+        return f"Invalid press type selected: {press_name}", -1, -1
 
     # Fetch the Wikipedia article
     try:
         content = requests.get(url).content
     except requests.exceptions.RequestException as e:
-        return f"Error fetching the Wikipedia article: {str(e)}", -1
+        return f"Error fetching the Wikipedia article: {str(e)}", -1, -1
 
     try:
         # Parse the Wikipedia HTML
@@ -48,12 +48,12 @@ def process_request(url, question, press_name, compression_ratio):
         num_tokens = pipe.tokenizer(context, return_tensors="pt")["input_ids"].shape[1]
         pred_answer = pipe(context, question=question, press=press)["answer"]
 
-        return pred_answer, num_tokens
+        return pred_answer, num_tokens, int(num_tokens * (1 - compression_ratio))
     except Exception as e:
         if "CUDA out of memory" in str(e):
             return "Error: CUDA out of memory. Try using a smaller article or a lower compression ratio.", -1
         else:
-            return str(e), -1
+            return str(e), -1, -1
 
 
 def gradio_interface():
@@ -82,10 +82,11 @@ def gradio_interface():
                 value="ExpectedAttentionPress",
                 label="Select Press Type",
             )
-            compression_slider = gr.Slider(minimum=0.1, maximum=1.0, step=0.1, value=0.5, label="Compression Ratio")
+            compression_slider = gr.Slider(minimum=0.0, maximum=0.9, step=0.1, value=0.5, label="Compression Ratio")
 
         output = gr.Textbox(label="Output", lines=10)
-        output_num_tokens = gr.Number(label="Number of Tokens", interactive=False)
+        output_num_tokens = gr.Number(label="Number of tokens before compression", interactive=False)
+        output_compressed_num_tokens = gr.Number(label="Number of tokens after compression", interactive=False)
 
         submit_button = gr.Button("Submit")
 
@@ -97,6 +98,12 @@ def gradio_interface():
                     "ExpectedAttentionPress",
                     0.5,
                 ],
+                [
+                    "https://en.wikipedia.org/wiki/Hugging_Face",
+                    "What was the original name of the transformers library ?",
+                    "ExpectedAttentionPress",
+                    0.5,
+                ],
             ],
             inputs=[url_input, question_input, press_selector, compression_slider],
         )
@@ -104,7 +111,7 @@ def gradio_interface():
         submit_button.click(
             process_request,
             inputs=[url_input, question_input, press_selector, compression_slider],
-            outputs=[output, output_num_tokens],
+            outputs=[output, output_num_tokens, output_compressed_num_tokens],
         )
 
     return demo