Update app.py

app.py

@@ -61,9 +61,7 @@ def process_text(text):
     text = re.sub(r'\[.*?\]', '', text, flags=re.DOTALL)
 
     return text
-
-import re
-
+    
 @spaces.GPU
 def generate(
     message: str,
@@ -85,42 +83,31 @@ def generate(
         gr.Warning(f"Trimmed input from conversation as it was longer than {MAX_INPUT_TOKEN_LENGTH} tokens.")
     
     streamer = TextIteratorStreamer(tokenizer, timeout=10.0, skip_prompt=True, skip_special_tokens=False)
-    generate_kwargs = {
-        "input_ids": input_ids,
-        "streamer": streamer,
-        "max_new_tokens": max_new_tokens,
-        "do_sample": True,
-        "top_p": top_p,
-        "top_k": top_k,
-        "temperature": temperature,
-        "num_beams": 1,
-        "repetition_penalty": repetition_penalty,
-    }
+    generate_kwargs = dict(
+        {"input_ids": input_ids},
+        streamer=streamer,
+        max_new_tokens=max_new_tokens,
+        do_sample=True,
+        top_p=top_p,
+        top_k=top_k,
+        temperature=temperature,
+        num_beams=1,
+        repetition_penalty=repetition_penalty,
+    )
     t = Thread(target=model.generate, kwargs=generate_kwargs)
     t.start()
 
     outputs = []
-    last_sentence_buffer = ""
     for text in streamer:
         processed_text = process_text(text)
-        sentences = re.split(r'(?<=\.)\s', processed_text)
-        
-        if len(sentences) > 1:
-            # Join all but the last sentence and buffer the last one
-            ready_to_stream = "".join(sentences[:-1])
-            if last_sentence_buffer:
-                yield last_sentence_buffer + ready_to_stream
-            last_sentence_buffer = sentences[-1]
-        else:
-            # No full sentences yet, buffer everything
-            last_sentence_buffer += processed_text
+        outputs.append(processed_text)
+        output = "".join(outputs)
+        yield output
 
     final_story = "".join(outputs)
-    if last_sentence_buffer:
-        final_story += last_sentence_buffer  # Add the last buffer if it's a complete sentence
-
-    # Optional: Save the final story without the last sentence
+    # Remove the last sentence
     final_story_trimmed = remove_last_sentence(final_story)
+
     try:
         saved_story = Story(message=message, content=final_story_trimmed).save()
         yield f"{final_story_trimmed}\n\n Story saved with ID: {saved_story.story_id}"
@@ -128,7 +115,7 @@ def generate(
         yield f"Failed to save story: {str(e)}"
 
 def remove_last_sentence(text):
-    # Split sentences and remove the last one
+    # Assuming sentences end with a period followed by space or end of string
     sentences = re.split(r'(?<=\.)\s', text)
     return ' '.join(sentences[:-1]) if sentences else text