Update app.py

app.py

@@ -7,6 +7,8 @@ from chromadb.utils import embedding_functions
 
 import gradio as gr
 
+import re 
+
 #######################################################
 
 # Load the email dataset
@@ -21,9 +23,9 @@ collection = client.create_collection("enron_emails")
 
 # Add documents and IDs to the collection, using ChromaDB's built-in text encoding
 collection.add(
-    documents=emails["body"].tolist()[:1000],
-    ids=emails["file"].tolist()[:1000],
-    metadatas=[{"source": "enron_emails"}] * len(emails[:1000]),  # Optional metadata
+    documents=emails["body"].tolist()[:10000],
+    ids=emails["file"].tolist()[:10000],
+    metadatas=[{"source": "enron_emails"}] * len(emails[:10000]),  # Optional metadata
 )
 
 
@@ -68,6 +70,10 @@ def summarize_documents(text_input):
         # Generate a summary with the model
         summary_ids = model.generate(inputs['input_ids'], max_length=150, min_length=40, length_penalty=2.0, num_beams=4, early_stopping=True)
         summary = tokenizer.decode(summary_ids[0], skip_special_tokens=True)
+
+        summary = re.sub(r"(\w+)([?!])\s", r"\1\2. ", summary)  # Ensures that sentences ending in ? ! .
+        summary = re.sub(r"([^.?!])(?=\s+[A-Z]|$)", r"\1.", summary) 
+        
         return summary
     except Exception as e:
         return f"An error occurred while summarizing: {e}"