Update app.py

app.py

@@ -39,7 +39,7 @@ st.sidebar.write("""
    - Click the Generate button.
 6. **Translation Tips:**
    - English as the target language gives the best results.
-   - You can also test inter-language translation i.e yoruba to igbo
+   - You can also do inter-language translation i.e yoruba to igbo
 7. **Performance Note:**
    - The model's performance varies due to its size and training data. It performs best on text generation and translation.
    - For other tasks, try multiple times if model's output is not optimal (This is due to the generator's sampling parameter settings).
@@ -115,6 +115,11 @@ def count_sentences(text):
     sentences = [sentence for sentence in sentences if sentence]
     return len(sentences)
 
+def wrap_text(text, task_value):
+    tasks = ["<classify>", "<prompt>", "<clean>", "<title>", "<diacritize>", "<translate>"]
+    if any(task in text for task in tasks):
+        return text
+    return task_value.format(text)
     
 async def generate_from_api(user_input, generation_config):
     urls = [
@@ -202,7 +207,8 @@ instruction_wrap = {
 #     "diacritize this text: ": "<diacritize> E sun, Alaga, fun ise amalayi ti e n se ni Naijiria. E maa ba a lo, egbon!  <yor>",
 #     "what is the topic of this text": "<classify> Africa Free Trade Zone: Kò sí ìdènà láti kó ọjà láti orílẹ̀èdè kan sí òmíràn <topic>",
 #     'headline of this text': '<title> ** Sylvain Itté French ambassador don comot Niger Republic **. Sylvain Itté, di French ambassador for Niger don comot Niamey and currently e dey for flight from Ndjamena to Paris. Sylvain Itté, di French ambassador for Niger don comot Niamey very early dis morning and currently e dey for flight from Ndjamena to Paris.\n\nDi military detain Bazoum and im family for di presidential palace. Niger na former French colony, and France still get 1,500 sojas for di African country.\n\n"France don decide to withdraw dia ambassador. In di next hours our ambassador and several diplomats go return to France," Oga Macron tok.\n\nE add say di military co-operation dey "over" and French troops go leave in "di months to come".\n\n"Dis Sunday we celebrate one new step towards di sovereignty of Niger," di junta tok, for one statement wey AFP news agency quote.\n\nDi decision by Paris dey come afta months of hostility and protest against di presence of French for di kontri, wit regular demonstrations for di capital Niamey.\n\nDi move don scata France operations against Islamist militants for di wider Sahel region and Paris influence for there. But oga Macron tok say "putschists no go hold France hostage,"  <headline>',
-# }
+# 
+
 # Task options
 task_options = {
     "select": "{}",
@@ -215,6 +221,7 @@ task_options = {
     "Text Diacritization": "<diacritize> {} ",
     "Question Generation": "<prompt> {} <response>:",
     "Question-Answering" : "<prompt> {} <response>:",
+    "Text Summarization" : "<summarize> {} <summary>:",
     "Text Cleaning": "<clean> {} "
 }
 # Language options for diacritize, translation and clean tasks
@@ -239,25 +246,24 @@ if task in ["Text Diacritization", "Text Cleaning", "Translation"]:
     language = st.selectbox("Select a Nigerian language:", list(language_options.keys()))
     task_value = f"{task_options[task]} {language_options[language]}"
 else:
-    task_value = task_options[task]
-def wrap_text(text, task_value):
-    tasks = ["<classify>", "<prompt>", "<clean>", "<title>", "<diacritize>", "<translate>"]
-    if any(task in text for task in tasks):
-        return text
-    return task_value.format(text)
+    task_value = task_options[task
+
+    
 # Text input 
 user_input = st.text_area("Enter text below **(PLEASE, FIRST READ ALL INSTRUCTIONS IN THE SIDEBAR CAREFULLY FOR THE BEST EXPERIENCE)**: ", sample_texts.get(sample_text, sample_text))
-n_sentences = count_sentences(user_input)
+
+if task == "Translation":
+    n_sentences = count_sentences(user_input)
 
 if task == "select":
     user_input = instruction_wrap.get(user_input, user_input)
                                      
-print("Final user input: ", user_input)
+# print("Final user input: ", user_input)
 if st.button("Generate"):
     if user_input:
         with st.spinner("Please wait..."):
             wrapped_input = wrap_text(user_input, task_value)
-            print("wrapped_input: ", wrapped_input)
+            # print("wrapped_input: ", wrapped_input)
             generation_config["max_new_tokens"]= min(max_new_tokens, 1024 - len(tokenizer.tokenize(wrapped_input)))
             start_time = time.time()
             
@@ -274,7 +280,7 @@ if st.button("Generate"):
             # generated_text  = re.split(r"\|(end_f_text|end_of_text|end_ofext|end_of_text_||end_of_te|end_oftext)|:|`", generated_text)[0]
             generated_text  = re.sub(r"\|(end_f_text|end_of_text|end_ofext|end_of_text_|end_of_te|end_o|end_of_tet|end_oftext)|:|`\|", "", generated_text)
             generated_text = generated_text.strip("\n")
-            print("Generated text: ", generated_text)
+            # print("Generated text: ", generated_text)
             
             if task == "Sentiment Classification" or "<sentiment>" in wrapped_input :
                 if "negative" in generated_text.lower():
@@ -286,12 +292,12 @@ if st.button("Generate"):
     
             elif task == "Topic Classification" or "<topic>" in wrapped_input:
                 generated_text  = generated_text[:15]
-                print("split", generated_text.split(" ")[0],  re.split(r"\.|\n|\*\*|\*", generated_text)[0], generated_text.split(" "))
+                # print("split", generated_text.split(" ")[0],  re.split(r"\.|\n|\*\*|\*", generated_text)[0], generated_text.split(" "))
                 generated_text = re.split(r"\.|\n|\*\*|\*", generated_text)[0]
                 generated_text = asyncio.run(assign_topic(generated_text))
                 
             elif task == "Translation" or "<translate>" in wrapped_input:
-                print("split for translation: ", n_sentences, re.split(r"\.|\n", generated_text)[:n_sentences])
+                # print("split for translation: ", n_sentences, re.split(r"\.|\n", generated_text)[:n_sentences])
                 generated_text = ".".join(re.split(r"\.|\n", generated_text)[:n_sentences])
 
             elif task == "Question Generation" or "Question Generation:" in sample_text: