updated app

app.py

@@ -5,6 +5,7 @@ from dotenv import load_dotenv
 import logging
 import sys  # Ensure sys is imported
 from huggingface_hub import login, HfApi
+import torch
 
 # Load environment variables
 load_dotenv()
@@ -51,41 +52,42 @@ except Exception as e:
 # Function to generate a formatted email
 def generate_email(recipient_name, recipient_email, industry, recipient_role, details):
     prompt = (
-        f"Write a cold outreach email for a {recipient_role} "
-        f"working in the {industry} industry. Use the following details: {details}. "
-        "The email should be professional and engaging, without any additional instructions or templates."
+        f"Write a short professional email to {recipient_name}, "
+        f"a {recipient_role} in the {industry} industry. "
+        f"Mention: {details}. Keep it under 100 words."
     )
-    inputs = tokenizer(prompt, return_tensors="pt")
-    outputs = model.generate(**inputs, max_length=300, num_return_sequences=1, temperature=0.7)
-    email_body = tokenizer.decode(outputs[0], skip_special_tokens=True)
-
-    # Extract only the relevant email content
-    email_lines = email_body.split('\n')
-    relevant_lines = []
-    for line in email_lines:
-        if line.strip().lower().startswith(('dear', 'hello', 'hi')):
-            relevant_lines = [line]
-        elif relevant_lines:
-            if line.strip().lower().startswith(('sincerely', 'best regards', 'regards', 'thank you')):
-                break
-            relevant_lines.append(line)
-
-    cleaned_email_body = '\n'.join(relevant_lines).strip()
-
-    # Format the email
-    formatted_email = f"""\
-To: {recipient_name} <{recipient_email}>
-Subject: Collaboration Opportunity
+    try:
+        inputs = tokenizer(prompt, return_tensors="pt", padding=True, truncation=True, max_length=512)
+        
+        # Check if we're using CPU or GPU
+        device = "cuda" if torch.cuda.is_available() else "cpu"
+        model.to(device)
+        inputs = {k: v.to(device) for k, v in inputs.items()}
+        
+        with torch.no_grad():
+            outputs = model.generate(
+                **inputs, 
+                max_length=200,  # Reduced max length
+                num_return_sequences=1,
+                temperature=0.7,
+                top_k=50,
+                top_p=0.95,
+                do_sample=True
+            )
+        
+        email_body = tokenizer.decode(outputs[0], skip_special_tokens=True)
+        logging.info(f"Generated raw email body: {email_body}")
 
-{cleaned_email_body}
+        # Minimal cleaning
+        email_body = email_body.strip()
+        if not email_body:
+            raise ValueError("Generated email body is empty")
 
-Best regards,
-Jane Smith
-Android Developer
-Albertsons
-[Your Contact Information]
+        # Format the email
+        formatted_email = f"""\
+To: {recipient_name} 
 """
-    return formatted_email
+        return formatted_email
 
 # Create Gradio interface
 iface = gr.Interface(