Update app.py

app.py

@@ -102,7 +102,7 @@ popular_topics = [
     "Philosophy", "Religion", "Society", "World"
 ]
 
-def extract_answer(text):
+async def extract_answer(text):
     pattern  = r'[a-z][A-Z]'
     result = re.split(pattern, text)[0]
     result = text[:len(result) + 1]
@@ -115,14 +115,14 @@ async def assign_topic(generated_text, topic_list=popular_topics):
             return topic.title()
     return generated_text
 
-def count_sentences(text):
+async def count_sentences(text):
     # Split the text based on sentence-ending punctuation followed by a space or end of string
     sentences = re.split(r'[.!?]+\s*', text.strip())
     # Filter out any empty strings from the resulting list
     sentences = [sentence for sentence in sentences if sentence]
     return len(sentences), sentences
 
-def wrap_text(text, task_value):
+async def wrap_text(text, task_value):
     tasks = ["<classify>", "<prompt>", "<clean>", "<title>", "<diacritize>", "<translate>"]
     if any(task in text for task in tasks):
         return text
@@ -271,7 +271,7 @@ else:
 if st.button("Generate"):
     if user_input:
         with st.spinner("Please wait..."):
-            wrapped_input = wrap_text(user_input, task_value)
+            wrapped_input = asyncio.run(wrap_text(user_input, task_value))
             # print("wrapped_input: ", wrapped_input)
             generation_config["max_new_tokens"]= min(max_new_tokens, 1024 - len(tokenizer.tokenize(wrapped_input)))
             start_time = time.time()
@@ -307,14 +307,17 @@ if st.button("Generate"):
                 
             elif task == "Translation" or "<translate>" in wrapped_input:
                 # print("split for translation: ", n_sentences, re.split(r"\.|\n", generated_text)[:n_sentences])
-                n_sentences, split_= count_sentences(initial_input)
+                n_sentences, split_= asyncio.run(count_sentences(initial_input))
                 print(n_sentences, split_)
-                _, generated_text = count_sentences(generated_text)
+                _, generated_text = asyncio.run(count_sentences(generated_text))
                 generated_text = ". ".join(generated_text[:n_sentences]) + "."
 
             elif task == "Question Generation" or "Question Generation:" in sample_text:
                 if "?" in generated_text:
                     generated_text = "? ".join(re.split(r"\?", generated_text)[:-1]) + "?"
+
+            elif task == "Question-Answering" or "Question-Answering:" in sample_text:
+                generated_text = asyncio.run(extract_answer(generated_text)
                     
                     
             full_output = st.empty()