updated the app

app.py

@@ -3,18 +3,19 @@ from transformers import AutoTokenizer, AutoModelForCausalLM, LlamaConfig
 import os
 from dotenv import load_dotenv
 import logging
-import sys  # Ensure sys is imported
+import sys
 from huggingface_hub import login, HfApi
 import torch
 
 # Load environment variables
 load_dotenv()
+
+# Setup logging
 logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s', stream=sys.stdout)
 
 # Authenticate with Hugging Face
 hf_token = os.environ.get("HUGGINGFACE_TOKEN")
 model_name = "meta-llama/Llama-3.1-8B"
-fallback_model = "facebook/opt-350m"
 
 if hf_token:
     try:
@@ -24,49 +25,107 @@ if hf_token:
         logging.info("Successfully logged in to Hugging Face")
     except Exception as e:
         logging.error(f"Error authenticating with Hugging Face: {str(e)}")
-        logging.warning("Proceeding without authentication. Will use fallback model.")
-        model_name = fallback_model
+        logging.warning("Proceeding without authentication. This may limit access to certain models.")
 else:
     logging.warning("HUGGINGFACE_TOKEN not found in environment variables. Proceeding without authentication.")
-    model_name = fallback_model
 
 # Load the model and tokenizer
 try:
+    logging.info(f"Attempting to load tokenizer for {model_name}")
     tokenizer = AutoTokenizer.from_pretrained(model_name)
+    logging.info("Tokenizer loaded successfully")
+    
+    logging.info(f"Attempting to load model {model_name}")
+    model = AutoModelForCausalLM.from_pretrained(model_name)
+    logging.info("Model loaded successfully")
     
-    # Custom configuration to handle the RoPE scaling issue
-    if model_name == "meta-llama/Llama-3.1-8B":
-        config = LlamaConfig.from_pretrained(model_name)
-        config.rope_scaling = {"type": "linear", "factor": 8.0}  # Adjust as needed
-        model = AutoModelForCausalLM.from_pretrained(model_name, config=config)
-    else:
-        model = AutoModelForCausalLM.from_pretrained(model_name)
+    if tokenizer.pad_token is None:
+        tokenizer.pad_token = tokenizer.eos_token
+        model.config.pad_token_id = model.config.eos_token_id
     
     logging.info(f"Successfully loaded {model_name}")
 except Exception as e:
     logging.error(f"Error loading {model_name}: {str(e)}")
-    logging.info(f"Falling back to {fallback_model}")
-    tokenizer = AutoTokenizer.from_pretrained(fallback_model)
-    model = AutoModelForCausalLM.from_pretrained(fallback_model)
+    raise
 
-def generate_email(recipient_name, recipient_email, industry, recipient_role, details):
-    prompt = (
-        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."
-    )
+MAX_TOTAL_TOKENS = 2048  # Adjusted to Llama model's token limit
+MAX_INPUT_TOKENS = 1600  # 1600 tokens for input, leaving room for generated output
+CONTEXT_RATIO = 0.6  # Adjusted for summarization
+
+def truncate_to_token_limit(text, max_tokens):
+    tokens = tokenizer.encode(text)
+    if len(tokens) > max_tokens:
+        tokens = tokens[:max_tokens]
+    return tokenizer.decode(tokens, skip_special_tokens=True)
+
+def summarize_text(text, max_tokens):
+    if len(tokenizer.encode(text)) <= max_tokens:
+        return text, None
+
+    summarization_prompt = f"""
+    Summarize the following text concisely, preserving the key points:
+
+    {text}
+
+    Ensure the summary is under {max_tokens} tokens.
+    """
+    
+    try:
+        inputs = tokenizer(summarization_prompt, return_tensors="pt", padding=True, truncation=True, max_length=int(MAX_INPUT_TOKENS * CONTEXT_RATIO))
+        
+        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():
+            summary_outputs = model.generate(
+                **inputs, 
+                max_new_tokens=max_tokens,
+                num_return_sequences=1,
+                temperature=0.7,
+                top_k=50,
+                top_p=0.95,
+                do_sample=True
+            )
+        
+        summary = tokenizer.decode(summary_outputs[0], skip_special_tokens=True)
+        summary = summary.replace(summarization_prompt, "").strip()
+        
+        warning = "Input was summarized to fit the token limit. Some details may be omitted."
+        return summary, warning
+    except Exception as e:
+        logging.error(f"Error during summarization: {str(e)}")
+        return text[:max_tokens] + "...", "Error in summarization. Text was truncated."
+
+def generate_prompt(recipient_name, recipient_role, industry, details):
+    details, warning = summarize_text(details, MAX_INPUT_TOKENS // 2)
+    
+    prompt_generation_input = f"""
+    Create a detailed prompt for writing a professional email based on the following information:
+    - Recipient: {recipient_name}, a {recipient_role} in the {industry} industry
+    - Purpose: {details}
+
+    Include:
+    1. Greeting
+    2. Main email points
+    3. Suggested closing
+    4. Tone (e.g., formal, friendly)
+    5. Industry-relevant phrases or terms
+    """
+
+    prompt_generation_input = truncate_to_token_limit(prompt_generation_input, MAX_INPUT_TOKENS)
+    
     try:
-        inputs = tokenizer(prompt, return_tensors="pt", padding=True, truncation=True, max_length=512)
+        inputs = tokenizer(prompt_generation_input, return_tensors="pt", padding=True, truncation=True, max_length=MAX_INPUT_TOKENS)
         
-        # 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(
+            prompt_outputs = model.generate(
                 **inputs, 
-                max_length=200,  # Reduced max length
+                max_new_tokens=200,
                 num_return_sequences=1,
                 temperature=0.7,
                 top_k=50,
@@ -74,34 +133,94 @@ def generate_email(recipient_name, recipient_email, industry, recipient_role, de
                 do_sample=True
             )
         
-        email_body = tokenizer.decode(outputs[0], skip_special_tokens=True)
-        logging.info(f"Generated raw email body: {email_body}")
+        generated_prompt = tokenizer.decode(prompt_outputs[0], skip_special_tokens=True)
+        return generated_prompt.replace(prompt_generation_input, "").strip(), warning
+
+    except Exception as e:
+        logging.error(f"Error generating prompt: {str(e)}")
+        return f"Error generating prompt: {str(e)}", None
+
+def generate_email_body(prompt):
+    # Concise prompt without instruction language
+    email_generation_input = f"""
+    {prompt}
+    """
+    
+    # Limit input to token constraints
+    email_generation_input = truncate_to_token_limit(email_generation_input, MAX_INPUT_TOKENS)
+    
+    try:
+        inputs = tokenizer(email_generation_input, return_tensors="pt", padding=True, truncation=True, max_length=MAX_INPUT_TOKENS)
+        
+        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():
+            email_outputs = model.generate(
+                **inputs, 
+                max_new_tokens=300,
+                num_return_sequences=1,
+                temperature=0.5,  # Lower temperature for more focused output
+                top_k=50,
+                top_p=0.95,
+                do_sample=False  # Deterministic output
+            )
+        
+        # Decode and return only the email body
+        email_body = tokenizer.decode(email_outputs[0], skip_special_tokens=True).strip()
+        
+        return email_body
 
-        # Minimal cleaning
-        email_body = email_body.strip()
-        if not email_body:
-            raise ValueError("Generated email body is empty")
+    except Exception as e:
+        logging.error(f"Error generating email body: {str(e)}")
+        return f"Error generating email body: {str(e)}"
 
-        # Format the email
-        formatted_email = f"""\
-To: {recipient_name} <{recipient_email}>
-Subject: Collaboration Opportunity
+def generate_email(recipient_name, recipient_email, industry, recipient_role, details):
+    try:
+        # Clear, minimal prompt to keep focus on main content generation
+        generated_prompt = f"I am reaching out to discuss {details} in the context of {industry} and how it impacts your role as a {recipient_role}."
+
+        # Generate the email body
+        email_body = generate_email_body(generated_prompt)
+        if email_body.startswith("Error"):
+            return email_body
+
+        # Remove duplicate greetings and signatures if they appear
+        email_body_lines = email_body.splitlines()
+        unique_lines = []
+        for line in email_body_lines:
+            if line.strip() and line not in unique_lines:
+                unique_lines.append(line)
+        email_body = "\n".join(unique_lines)
+
+        # Assemble the final email content
+        final_output = f"""
+Subject: {details.split()[0].capitalize()} Proposal
+
+Dear {recipient_name},
 
 {email_body}
 
-Best regards,
-Jane Smith
-Android Developer
-Albertsons
-[Your Contact Information]
+Sincerely,  
+Your Name  
+Your Title  
+Your Company  
+Your Email  
+Your Phone Number  
+Your Address  
+Your Website  
+Your Social Media Profiles
 """
-        return formatted_email
+        logging.info(f"Final email content:\n{final_output}")
+        return final_output
 
     except Exception as e:
-        logging.error(f"Error generating email: {str(e)}")
+        logging.error(f"Error in generate_email: {str(e)}")
         return f"Error generating email: {str(e)}"
 
-# Create Gradio interface
+
+
 iface = gr.Interface(
     fn=generate_email,
     inputs=[
@@ -109,17 +228,12 @@ iface = gr.Interface(
         gr.Textbox(lines=1, label="Recipient Email"),
         gr.Textbox(lines=1, label="Industry (e.g., Technology, Healthcare)"),
         gr.Textbox(lines=1, label="Recipient Role (e.g., Manager, Director)"),
-        gr.Textbox(lines=5, label="Personal/Company Details (e.g., name, product)"),
+        gr.Textbox(lines=5, label="Personal/Company Details"),
     ],
     outputs=gr.Textbox(lines=10, label="Generated Email or Error Message"),
     title="EmailGenie: AI-Powered Email Generator",
     description="Automate the creation of personalized emails. Enter details to generate tailored emails."
 )
 
-# Launch the app
 if __name__ == '__main__':
     iface.launch()
-
-# Log model information
-logging.info(f"Model name: {model.config._name_or_path}")
-logging.info(f"Model parameters: {model.num_parameters()}")