Imorove ui

app.py

@@ -8,12 +8,13 @@ def run_prompt(prompt, api_key, model_name, max_length):
     lemmatized_prompt = lemmatizer_func(prompt)
     response_plain = prompt_node(prompt)
     response_lemmatized = prompt_node(lemmatized_prompt)
-    return response_plain[0][0], response_plain[1]["total_tokens"], response_lemmatized[0][0], response_lemmatized[1]["total_tokens"]
+    return lemmatized_prompt, response_plain[0][0], response_plain[1]["prompt_tokens"], response_plain[1]["completion_tokens"], response_lemmatized[0][0], response_lemmatized[1]["prompt_tokens"], response_lemmatized[1]["completion_tokens"]
 
 with gr.Blocks() as demo:
     with gr.Row():
         api_key = gr.Textbox(label="Enter your api key")
         model_name = gr.Dropdown(["text-davinci-003", "gpt-3.5-turbo", "gpt-4", "gpt-4-32k", "command", "command-light", "base", "base-light"], value="gpt-3.5-turbo", label="Choose your model!")
+        max_length = gr.Slider(100, 500, value=100, step=10, label="Max Length", info="Max token length of the response. Choose between 100 and 500")   
     with gr.Row():
         prompt = gr.TextArea(label="Prompt", value="Rachel has 17 apples. She gives 9 to Sarah. How many apples does Rachel have now?")
         gr.Examples(
@@ -26,16 +27,20 @@ with gr.Blocks() as demo:
         ],
         inputs=prompt,
         label="Click on any example and press Enter in the input textbox!",
-    )
-    max_length = gr.Slider(100, 500, value=100, step=10, label="Max Length", info="Choose between 100 and 500")    
+    ) 
     submit_btn = gr.Button("Submit")
     with gr.Row():
         prompt_response = gr.TextArea(label="Answer", show_copy_button=True)
-        token_count_plain = gr.Number(label="Plain Text Token Count")
+        with gr.Column():
+            token_count_plain = gr.Number(label="Prompt Token Count")
+            token_count_plain_completion = gr.Number(label="Completion Token Count")
     with gr.Row():
         lemmatized_prompt_response = gr.TextArea(label="Lemmatized Answer", show_copy_button=True)
-        token_count_lemmatized = gr.Number(label="Lemmatized Text Token Count")
-    submit_btn.click(fn=run_prompt, inputs=[prompt, api_key, model_name, max_length], outputs=[prompt_response, token_count_plain, lemmatized_prompt_response, token_count_lemmatized])
+        with gr.Column():
+            lemmatized_prompt = gr.TextArea(label="Lemmatized Promp", show_copy_button=True)
+            token_count_lemmatized = gr.Number(label="Lemmatized Text Token Count")
+            token_count_lemmatized_completion = gr.Number(label="Completion Lemmatized Token Count")
+    submit_btn.click(fn=run_prompt, inputs=[prompt, api_key, model_name, max_length], outputs=[lemmatized_prompt, prompt_response, token_count_plain, token_count_plain_completion, lemmatized_prompt_response, token_count_lemmatized, token_count_lemmatized_completion])
 
 demo.launch()
 

utils.py

@@ -1,6 +1,5 @@
 from simplemma import text_lemmatizer
 
 def lemmatizer_func(plain_text):
-    words = text_lemmatizer(plain_text, lang="en")
-    lemmatized_promt = ' '.join(words)
+    lemmatized_promt = ' '.join([el for el in text_lemmatizer(plain_text, lang='en') if el.isalnum()])
     return lemmatized_promt