Update app.py

app.py

@@ -1,72 +1,64 @@
 import os
 import gradio as gr
-from transformers import AutoTokenizer, AutoModelForCausalLM, pipeline
-import torch
+from huggingface_hub import InferenceClient
 import logging
 
 logging.basicConfig(level=logging.INFO)
 logger = logging.getLogger(__name__)
 
-# Change to a more suitable model for text generation
-model_name = "meta-llama/Llama-3.2-1B"  # Ensure this model is available for your use
-
-# In Hugging Face Spaces, the token is usually available as an environment variable
+model_name = "meta-llama/Llama-3.2-1B"
 hf_token = os.environ.get("HUGGINGFACE_TOKEN")
 if not hf_token:
-    logger.warning("HUGGINGFACE_TOKEN environment variable is not set. Some features may not work.")
+    logger.error("HUGGINGFACE_TOKEN environment variable is not set")
+    raise ValueError("HUGGINGFACE_TOKEN environment variable is not set")
 
-try:
-    tokenizer = AutoTokenizer.from_pretrained(model_name, use_auth_token=hf_token)
-    model = AutoModelForCausalLM.from_pretrained(model_name, use_auth_token=hf_token)
-    generator = pipeline('text-generation', model=model, tokenizer=tokenizer)
-    logger.info(f"Successfully loaded model and tokenizer: {model_name}")
-except Exception as e:
-    logger.error(f"Failed to load model or tokenizer: {type(e).__name__}: {str(e)}")
-    # Fallback to a smaller, open-access model if the specified model fails to load
-    model_name = "distilgpt2"
-    tokenizer = AutoTokenizer.from_pretrained(model_name)
-    model = AutoModelForCausalLM.from_pretrained(model_name)
-    generator = pipeline('text-generation', model=model, tokenizer=tokenizer)
-    logger.info(f"Fallback: Loaded model and tokenizer: {model_name}")
+client = InferenceClient(model=model_name, token=hf_token)
 
 def generate_text(prompt):
     try:
         logger.info(f"Attempting to generate text for prompt: {prompt[:50]}...")
         
-        response = generator(prompt, max_length=1000, num_return_sequences=1, do_sample=True, temperature=0.7)
-        
-        generated_text = response[0]['generated_text']
+        response = client.text_generation(
+            prompt,
+            max_new_tokens=500,
+            temperature=0.7,
+            top_k=50,
+            top_p=0.95,
+            do_sample=True
+        )
         
-        logger.info(f"Generated text: {generated_text[:100]}...")
-        return generated_text
+        logger.info(f"Generated text: {response[:100]}...")
+        return response
     except Exception as e:
         logger.error(f"Error in generate_text: {type(e).__name__}: {str(e)}")
         return f"An error occurred: {type(e).__name__}: {str(e)}"
 
 def generate_email(industry, recipient_role, company_details):
     try:
-        prompt = f"""Generate a professional cold outreach email using the following details:
-        Industry: {industry}
-        Recipient Role: {recipient_role}
-        Company Details: {company_details}
-        Format the email as follows:
-        Subject: [Insert catchy subject line related to the industry and recipient role]
-        Dear [Recipient's Name],
-        [Introduction paragraph: Briefly introduce yourself and your company]
-        [Value proposition paragraph: Explain how your company can benefit the recipient, citing specific details from the company information]
-        [Call to action paragraph: Suggest a meeting or call to discuss further]
-        [Closing paragraph: Thank the recipient and provide your contact information]
-        Best regards,
-        [Your Name]
-        [Your Title]
-        [Your Company]
-        Email:
-        """
+        prompt = f"""Task: Generate a professional cold outreach email.
+
+Context:
+- Industry: {industry}
+- Recipient Role: {recipient_role}
+- Company Details: {company_details}
+
+Instructions:
+1. Create a catchy subject line related to the industry and recipient role.
+2. Write a personalized greeting.
+3. Introduce yourself and your company briefly.
+4. Explain how your company can benefit the recipient, using specific details from the company information.
+5. Suggest a meeting or call to discuss further.
+6. Thank the recipient and provide your contact information.
+7. Use a professional closing.
+
+Now, write the email following these instructions. Be creative and specific, don't use placeholder text:
+
+"""
         
         generated_text = generate_text(prompt)
         
-        # Remove the prompt from the generated text
-        email_content = generated_text[generated_text.find("Subject:"):].strip()
+        # Remove any remaining prompt text if present
+        email_content = generated_text.split("Now, write the email following these instructions.")[-1].strip()
         
         logger.info(f"Generated email for {industry}, {recipient_role}")
         return email_content
@@ -76,7 +68,7 @@ def generate_email(industry, recipient_role, company_details):
 
 def test_model_connection():
     try:
-        test_prompt = "Generate a short, coherent paragraph about artificial intelligence:"
+        test_prompt = "Write a short paragraph about the importance of AI in modern business:"
         response = generate_text(test_prompt)
         logger.info(f"Test model connection successful. Response: {response}")
         return "Model connection test successful. Response: " + response