app.py

import gradio as gr
from transformers import AutoTokenizer, AutoModelForCausalLM, LlamaConfig
import os
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()
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:
        login(token=hf_token)
        api = HfApi()
        api.whoami()
        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
else:
    logging.warning("HUGGINGFACE_TOKEN not found in environment variables. Proceeding without authentication.")
    model_name = fallback_model

# Load the model and tokenizer
try:
    tokenizer = AutoTokenizer.from_pretrained(model_name)
    
    # 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)
    
    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)

# Function to generate a formatted email
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."
    )
    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}")

        # Minimal cleaning
        email_body = email_body.strip()
        if not email_body:
            raise ValueError("Generated email body is empty")

        # Format the email
        formatted_email = f"""\
To: {recipient_name} 
"""
        return formatted_email

# Create Gradio interface
iface = gr.Interface(
    fn=generate_email,
    inputs=[
        gr.Textbox(lines=1, label="Recipient Name"),
        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)"),
    ],
    outputs="text",
    title="EmailGenie: AI-Powered Email Generator",
    description="Automate the creation of personalized emails to increase engagement and conversion rates. Enter details to generate tailored emails."
)

# Launch the app
if __name__ == '__main__':
    iface.launch()