Update app.py

app.py

@@ -62,7 +62,8 @@ def process_text(text):
 
     return text
 
-# Gradio Function
+import re
+
 @spaces.GPU
 def generate(
     message: str,
@@ -84,34 +85,54 @@ 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 = 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,
-    )
+    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,
+    }
     t = Thread(target=model.generate, kwargs=generate_kwargs)
     t.start()
 
     outputs = []
+    last_sentence_buffer = ""
     for text in streamer:
         processed_text = process_text(text)
-        outputs.append(processed_text)
-        output = "".join(outputs)
-        yield output
+        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
 
     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
+    final_story_trimmed = remove_last_sentence(final_story)
     try:
-        saved_story = Story(message=message, content=final_story).save()
-        yield f"{final_story}\n\n Story saved with ID: {saved_story.story_id}"
+        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}"
     except Exception as e:
         yield f"Failed to save story: {str(e)}"
 
+def remove_last_sentence(text):
+    # Split sentences and remove the last one
+    sentences = re.split(r'(?<=\.)\s', text)
+    return ' '.join(sentences[:-1]) if sentences else text
+
+
 # Gradio Interface Setup
 chat_interface = gr.ChatInterface(
     fn=generate,