Update utils.py

utils.py

@@ -10,31 +10,42 @@ model = AutoModelForCausalLM.from_pretrained(model_checkpoint)
 generator = pipeline('text-generation', model=model, tokenizer=tokenizer)
 
 def clean_generated_text(text):
-    # Basic cleaning
-    text = re.sub(r'^(Re:|Fwd:)', '', text)  # Remove reply and forward marks
-    text = re.sub(r'Best regards,.*$', '', text, flags=re.DOTALL)  # Remove everything after signature
-    text = re.sub(r'PHONE.*$', '', text, flags=re.DOTALL)  # Remove everything after phone numbers
-    text = re.sub(r'Email:.*$', '', text, flags=re.DOTALL)  # Remove everything after email addresses
-    text = re.sub(r'Cc:.*$', '', text, flags=re.DOTALL)  # Remove CC list
-    text = re.sub(r'\* Attachments:.*', '', text, flags=re.S)  # Remove 'Attachments:' and everything following it
-    text = re.sub(r'©️ .*$', '', text, flags=re.DOTALL)  # Remove copyright and ownership statements
-    text = re.sub(r'URL If this message is not displaying properly, click here.*$', '', text, flags=re.DOTALL)  # Remove error display message and links
-    text = re.sub(r'\d{5,}', 'NUMBER', text)  # Replace long sequences of numbers, likely phone numbers or ZIP codes
+    #Basic cleaning
+    text = re.sub(r'^(Re:|Fwd:)', '', text)                         # Remove reply and forward marks
+    text = re.sub(r'Best regards,.*$', '', text, flags=re.DOTALL)   # Remove signature
+    text = re.sub(r'PHONE.*$', '', text, flags=re.DOTALL)           # Remove phone numbers
+    text = re.sub(r'Email:.*$', '', text, flags=re.DOTALL)          # Remove email addresses
+    text = re.sub(r'Cc:.*$', '', text, flags=re.DOTALL)             # Remove CC list
+    text = re.sub(r'\* Attachments:.*', '', text, flags=re.S)       # Remove Attachments
+    text = re.sub(r'©️ .*$', '', text, flags=re.DOTALL)              # Remove copyright and ownership statements
+    text = re.sub(r'URL', '', text)                                 # Remove URLs
+    text = re.sub(r'NUMBER', '10', text)                            # Replace 'NUMBER' with a real number
+    text = re.sub(r'CURRENCYNUMBER', 'USD 100', text)               # Replace 'CURRENCYNUMBER' with a real value
+    text = re.sub(r'About Us.*', '', text, flags=re.DOTALL)         # Remove 'About Us' and all following text
+    text = re.sub(r'\d+ [^\s]+ St\.?,?.*?\d{5}', '', text)          # Remove street
+    text = re.sub(r'\d+ [^\s]+ Ave\.?,?.*?\d{5}', '', text)         # Remove avenues
+    text = re.sub(r'\d+ [^\s]+ Rd\.?,?.*?\d{5}', '', text)          # Remove roads
+    text = re.sub(r'\d+ [^\s]+ Ln\.?,?.*?\d{5}', '', text)          # Remove lanes
+    text = re.sub(r'\d+ [^\s]+ Blvd\.?,?.*?\d{5}', '', text)        # Remove boulevards
+    text = re.sub(r'\d+ [^\s]+ Dr\.?,?.*?\d{5}', '', text)          # Remove drives
+    text = re.sub(r'\d+ [^\s]+ Ct\.?,?.*?\d{5}', '', text)          # Remove courts
     return text.strip()
 
 def generate_email(product, gender, profession, hobby):
     input_text = f"{product} {gender} {profession} {hobby}"
     result = generator(
-        input_text,                 # Initial text to prompt the model. Sets the context or topic for text generation.
-        max_length=256,             # Maximum length of the generated text in tokens, limiting the output size.
-        do_sample=True,             # Enables stochastic sampling; the model can generate diverse outputs at each step.
-        top_k=20,                   # Limits the vocabulary considered at each step to the top-k most likely next words.
-        top_p=0.6,                  # Uses nucleus sampling: Narrows down to the smallest set of words totaling 60% of the likelihood.
-        temperature=0.4,            # Scales logits before sampling to reduce randomness and produce more deterministic output.
-        repetition_penalty=1.5,     # Penalizes words that were already mentioned, reducing repetition in the text.
-        # truncation=True,            # Truncates the output to the maximum length if it exceeds it.
-        num_return_sequences=3      # Generates three different sequences to choose from, enhancing output variety.
-        ) 
-    # Select the best output from the generated sequences
-    best_text = sorted([clean_generated_text(r['generated_text']) for r in result], key=len)[-1]
+        input_text,                # The starting text that guides the model on what to generate
+        max_length=256,            # Set a suitable maximum length
+        top_k=40,                  # Consider more top options words
+        top_p=0.6,                 # Control the probability range for word choices
+        temperature=0.4,           # Control the randomness of generation
+        repetition_penalty=1.5,    # Reduce content repetition
+        num_return_sequences=2,    # Generate three texts
+        do_sample=True
+    )       
+    # Clean each generated text
+    cleaned_texts = [clean_generated_text(seq['generated_text']) for seq in result]
+    # Choose the best text based on length and clarity
+    best_text = max(cleaned_texts, key=len)
     return best_text
+