app.py

import os
import gradio as gr
from huggingface_hub import InferenceClient
import logging

# Set up logging
logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)

# Get the token from the environment variable set in HF Spaces
hf_token = os.environ.get("HUGGINGFACE_TOKEN")
if not hf_token:
    logger.error("HUGGINGFACE_TOKEN environment variable is not set")
    raise ValueError("HUGGINGFACE_TOKEN environment variable is not set")

# Model name
model_name = "meta-llama/Llama-2-7b-chat-hf"

# Create an InferenceClient
client = InferenceClient(model=model_name, token=hf_token)

def generate_text(prompt):
    try:
        logger.info(f"Generating text for prompt: {prompt[:50]}...")
        response = client.text_generation(prompt, max_new_tokens=500, temperature=0.7)
        return response
    except Exception as e:
        logger.error(f"Error in generate_text: {str(e)}")
        return f"An error occurred: {str(e)}"

def generate_email(industry, recipient_role, company_details):
    try:
        logger.info(f"Generating email for {industry}, {recipient_role}")
        prompt = f"Generate a personalized cold outreach email for a {recipient_role} in the {industry} industry. Company details: {company_details}"
        return generate_text(prompt)
    except Exception as e:
        logger.error(f"Error in generate_email: {str(e)}")
        return f"An error occurred in generate_email: {str(e)}"

def refine_email(email, tone, context):
    try:
        logger.info(f"Refining email with tone: {tone}")
        prompt = f"Refine the following email to match a {tone} tone and consider this context: {context}\n\nEmail: {email}"
        return generate_text(prompt)
    except Exception as e:
        logger.error(f"Error in refine_email: {str(e)}")
        return f"An error occurred in refine_email: {str(e)}"

# Gradio interface
with gr.Blocks() as demo:
    gr.Markdown("# EmailGenie: AI-Powered Cold Email Generator")
    
    with gr.Tab("Generate Email"):
        industry = gr.Textbox(label="Industry")
        recipient_role = gr.Textbox(label="Recipient Role")
        company_details = gr.Textbox(label="Company/Personal Details", lines=5)
        generate_button = gr.Button("Generate Email")
        output = gr.Textbox(label="Generated Email", lines=10)
        
        generate_button.click(generate_email, inputs=[industry, recipient_role, company_details], outputs=output)
    
    with gr.Tab("Refine Email"):
        email_input = gr.Textbox(label="Input Email", lines=10)
        tone = gr.Dropdown(["Formal", "Casual", "Enthusiastic"], label="Desired Tone")
        context = gr.Textbox(label="Additional Context", lines=3)
        refine_button = gr.Button("Refine Email")
        refined_output = gr.Textbox(label="Refined Email", lines=10)
        
        refine_button.click(refine_email, inputs=[email_input, tone, context], outputs=refined_output)

demo.launch()