utils

utils.py

@@ -0,0 +1,40 @@
+from transformers import pipeline, AutoModelForCausalLM, AutoTokenizer
+import re
+
+# Load saved model and tokenizer
+model_checkpoint = "24NLPGroupO/EmailGeneration"
+tokenizer = AutoTokenizer.from_pretrained(model_checkpoint, truncation=True)
+model = AutoModelForCausalLM.from_pretrained(model_checkpoint)
+
+# Set up the generation pipeline
+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
+    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]
+    return best_text