Update app.py

app.py

@@ -2,62 +2,225 @@ import gradio as gr
 import uuid
 from datetime import datetime
 import random
+import os
+from transformers import pipeline, AutoTokenizer, AutoModelForCausalLM
+import torch
+from huggingface_hub import login, HfApi
 
-# Mock function to simulate AI generation
-def generate_content(prompt, max_length=150):
-    """Simulate AI content generation"""
-    topics = prompt.lower().split()
+# Token setup - Read from environment variables
+HF_TOKEN_READ = os.environ.get("HF_TOKEN_READ", "")
+HF_TOKEN_WRITE = os.environ.get("HF_TOKEN_WRITE", "")
+
+# Login with appropriate token
+if HF_TOKEN_READ:
+    try:
+        login(token=HF_TOKEN_READ)
+        print("Logged in with read token")
+    except Exception as e:
+        print(f"Error logging in with read token: {e}")
+
+# Initialize HF API for write operations (if token is available)
+hf_api = None
+if HF_TOKEN_WRITE:
+    try:
+        hf_api = HfApi(token=HF_TOKEN_WRITE)
+        print("HF API initialized with write token")
+    except Exception as e:
+        print(f"Error initializing HF API: {e}")
+
+# Model loading function with token support
+@gr.Cache()
+def load_model(model_name="microsoft/DialoGPT-medium"):
+    """Load a Hugging Face model for text generation"""
+    try:
+        # Use read token if available
+        token = HF_TOKEN_READ if HF_TOKEN_READ else None
+        
+        tokenizer = AutoTokenizer.from_pretrained(
+            model_name, 
+            token=token,
+            use_auth_token=token is not None
+        )
+        
+        if tokenizer.pad_token is None:
+            tokenizer.pad_token = tokenizer.eos_token
+            
+        model = AutoModelForCausalLM.from_pretrained(
+            model_name,
+            token=token,
+            use_auth_token=token is not None
+        )
+        
+        return tokenizer, model
+    except Exception as e:
+        print(f"Error loading model: {e}")
+        return None, None
+
+# Save content to Hugging Face Hub (if write token is available)
+def save_to_hub(content, filename="social_content.txt", repo_id="your-username/social-media-content"):
+    """Save generated content to Hugging Face Hub"""
+    if not hf_api or not HF_TOKEN_WRITE:
+        return "Write token not configured. Content not saved to Hub."
+    
+    try:
+        # Create a temporary file
+        with open(filename, "w") as f:
+            f.write(content)
+        
+        # Upload to Hub
+        hf_api.upload_file(
+            path_or_fileobj=filename,
+            path_in_repo=filename,
+            repo_id=repo_id,
+            repo_type="dataset",
+            commit_message=f"Add social media content - {datetime.now().strftime('%Y-%m-%d %H:%M')}"
+        )
+        
+        # Clean up
+        os.remove(filename)
+        
+        return f"Content saved to Hub: https://huggingface.co/datasets/{repo_id}"
+    except Exception as e:
+        return f"Error saving to Hub: {e}"
+
+# Fallback content generation
+def fallback_generate_content(prompt, platform):
+    """Fallback content generation if AI models fail"""
     responses = {
-        "instagram": [
+        "Instagram": [
             f"✨ {prompt.title()} ✨\n\nReady to transform your approach? Here's what you need to know:\n\n🔥 Key insight that changes everything\n💡 Pro tip that most people miss\n⚡ Action step you can take today\n\nWhat's your experience? Share in the comments! 👇",
             f"🌟 Behind the scenes of {prompt} 🌟\n\nSharing our process and what makes it special!\n\nEver wondered about {prompt}? Let us know your questions below! ⬇️",
-            f"📸 Capturing the essence of {prompt}\n\nSometimes the simplest moments become the most memorable. \n\nTag someone who needs to see this! 🏷️"
         ],
-        "tiktok": [
+        "TikTok": [
             f"POV: You finally understand {prompt} 🤯\n\n*shows before and after*\n\nThe secret? [key insight]\n\nWho else needed to hear this? 💪",
             f"Wait until you try this {prompt} hack! 👀\n\nGame changer alert! 🚨\n\nSave this for later! ⬇️",
-            f"Duet this if you agree! 👇\n\n{prompt} is seriously underrated. \n\nWhat's your take? 🤔"
         ],
-        "both": [
+        "Both": [
             f"Cross-platform content for {prompt} 📱\n\nCreating value across different channels!\n\nWhat platform do you prefer? Let me know! 🗣️",
-            f"Multi-platform strategy for {prompt} 🎯\n\nDifferent content for different audiences!\n\nWhich platform works best for you? 📊"
         ]
     }
     
-    # Simple logic to select response based on platform hint in prompt
-    if "instagram" in prompt.lower():
-        return random.choice(responses["instagram"])
-    elif "tiktok" in prompt.lower():
-        return random.choice(responses["tiktok"])
+    return random.choice(responses.get(platform, responses["Both"]))
+
+# AI-powered content generation
+def generate_content(prompt, platform, max_length=150):
+    """Generate content using Hugging Face models"""
+    try:
+        # Load an appropriate model based on platform
+        if platform == "Instagram":
+            model_name = "microsoft/DialoGPT-medium"
+        else:
+            model_name = "gpt2"
+        
+        tokenizer, model = load_model(model_name)
+        
+        if tokenizer is None or model is None:
+            return fallback_generate_content(prompt, platform)
+        
+        # Format prompt for better results
+        formatted_prompt = f"Create engaging {platform} content about: {prompt}"
+        
+        # Generate content
+        inputs = tokenizer.encode(formatted_prompt, return_tensors="pt")
+        attention_mask = torch.ones(inputs.shape, dtype=torch.long)
+        
+        outputs = model.generate(
+            inputs, 
+            max_length=max_length, 
+            num_return_sequences=1,
+            temperature=0.8,
+            do_sample=True,
+            pad_token_id=tokenizer.eos_token_id,
+            attention_mask=attention_mask
+        )
+        
+        generated_text = tokenizer.decode(outputs[0], skip_special_tokens=True)
+        
+        # Clean up the output
+        if generated_text.startswith(formatted_prompt):
+            generated_text = generated_text[len(formatted_prompt):].strip()
+        
+        return generated_text if generated_text else fallback_generate_content(prompt, platform)
+        
+    except Exception as e:
+        print(f"Error in generate_content: {e}")
+        return fallback_generate_content(prompt, platform)
+
+# Improved hashtag generation
+def generate_hashtags(topic, platform):
+    """Generate relevant hashtags using AI"""
+    try:
+        # Use read token if available
+        token = HF_TOKEN_READ if HF_TOKEN_READ else None
+        
+        # Use a pipeline for text generation
+        generator = pipeline(
+            'text-generation', 
+            model='gpt2',
+            token=token,
+            use_auth_token=token is not None
+        )
+        
+        prompt = f"Generate 5 relevant hashtags for {topic} on {platform}:"
+        result = generator(prompt, max_length=50, num_return_sequences=1)
+        
+        hashtags = result[0]['generated_text'].replace(prompt, '').strip()
+        # Clean up and format hashtags
+        hashtag_list = [tag.strip().replace(' ', '') for tag in hashtags.split()[:5]]
+        hashtags = ' '.join(['#' + tag for tag in hashtag_list if tag])
+        
+        return hashtags if hashtags else fallback_hashtags(topic, platform)
+        
+    except Exception as e:
+        print(f"Error generating hashtags: {e}")
+        return fallback_hashtags(topic, platform)
+
+def fallback_hashtags(topic, platform):
+    """Fallback hashtag generation"""
+    base_hashtags = topic.lower().replace(' ', '').replace(',', ' #')
+    
+    if platform == "Instagram":
+        platform_tags = "#instagood #photooftheday #instadaily #motivation #inspiration"
+    elif platform == "TikTok":
+        platform_tags = "#fyp #foryou #viral #trending #tiktok"
     else:
-        return random.choice(responses["both"])
+        platform_tags = "#content #socialmedia #digital"
+    
+    return f"#{base_hashtags} {platform_tags}"
+
+# Content improvement function
+def improve_content(content, platform, tone):
+    """Improve existing content using AI"""
+    try:
+        # Use read token if available
+        token = HF_TOKEN_READ if HF_TOKEN_READ else None
+        
+        improver = pipeline(
+            'text2text-generation', 
+            model='google/flan-t5-base',
+            token=token,
+            use_auth_token=token is not None
+        )
+        
+        prompt = f"Improve this {platform} content to make it more {tone}: {content}"
+        result = improver(prompt, max_length=200)
+        
+        return result[0]['generated_text']
+    except Exception as e:
+        print(f"Error improving content: {e}")
+        return content
 
 # Main content generation function
-def generate_social_media_content(platform, content_type, topic, target_audience, tone, length, brand_voice, key_message, call_to_action):
+def generate_social_media_content(platform, content_type, topic, target_audience, tone, length, brand_voice, key_message, call_to_action, model_choice):
     """Generate social media content based on inputs"""
     
     # Generate main content
     content_prompt = f"Create a {content_type.lower()} for {platform} about {topic} targeting {target_audience} in a {tone.lower()} tone"
-    generated_content = generate_content(content_prompt)
+    generated_content = generate_content(content_prompt, platform)
     
     # Generate hashtags
     hashtags = generate_hashtags(topic, platform)
     
-    # Create content object
-    new_content = {
-        'id': str(uuid.uuid4()),
-        'platform': platform,
-        'type': content_type,
-        'topic': topic,
-        'content': generated_content,
-        'hashtags': hashtags,
-        'audience': target_audience,
-        'tone': tone,
-        'cta': call_to_action,
-        'created': datetime.now().strftime('%Y-%m-%d %H:%M'),
-        'status': 'Draft'
-    }
-    
     # Calculate stats
     content_text = generated_content
     char_count = len(content_text)
@@ -95,35 +258,18 @@ def generate_social_media_content(platform, content_type, topic, target_audience
 - Target: {target_audience}
 - Tone: {tone}
 - CTA: {call_to_action}
-- Created: {new_content['created']}
+- AI Model: {model_choice}
+- Created: {datetime.now().strftime('%Y-%m-%d %H:%M')}
 """
     
     return output
 
-def generate_hashtags(topic, platform):
-    """Generate relevant hashtags"""
-    base_hashtags = topic.lower().replace(' ', '').replace(',', ' #')
-    
-    if platform == "Instagram":
-        platform_tags = "#instagood #photooftheday #instadaily #motivation #inspiration"
-    elif platform == "TikTok":
-        platform_tags = "#fyp #foryou #viral #trending #tiktok"
-    else:
-        platform_tags = "#content #socialmedia #digital"
-    
-    return f"#{base_hashtags} {platform_tags}"
-
 def get_trending_topics():
     """Get trending topic suggestions"""
     trending = [
-        "Productivity hacks",
-        "Morning routine", 
-        "Self care Sunday",
-        "Workspace setup",
-        "Healthy recipes",
-        "Weekend vibes",
-        "Goal setting",
-        "Mindfulness tips"
+        "Productivity hacks", "Morning routine", "Self care Sunday",
+        "Workspace setup", "Healthy recipes", "Weekend vibes",
+        "Goal setting", "Mindfulness tips"
     ]
     return random.choice(trending)
 
@@ -168,20 +314,64 @@ with gr.Blocks(
         padding: 20px;
         border-radius: 10px;
         background: #f8f9fa;
+        border-left: 4px solid #667eea;
+    }
+    .header {
+        text-align: center;
+        padding: 20px;
+        background: linear-gradient(135deg, #667eea 0%, #764ba2 100%);
+        color: white;
+        border-radius: 10px;
+        margin-bottom: 20px;
+    }
+    .token-status {
+        padding: 10px;
+        border-radius: 5px;
+        margin-bottom: 10px;
+    }
+    .token-ok {
+        background: #d4edda;
+        color: #155724;
+        border: 1px solid #c3e6cb;
+    }
+    .token-warning {
+        background: #fff3cd;
+        color: #856404;
+        border: 1px solid #ffeeba;
     }
     """
 ) as demo:
-    gr.Markdown("# 📱 Social Media Content Creator")
-    gr.Markdown("Create engaging content for Instagram and TikTok with AI assistance")
+    with gr.Column():
+        gr.Markdown("""
+        <div class="header">
+            <h1>📱 Social Media Content Creator</h1>
+            <p>Create engaging content for Instagram and TikTok with AI assistance</p>
+        </div>
+        """)
+    
+    # Token status indicator
+    token_status = "🔑 Read Token: " + ("✅ Available" if HF_TOKEN_READ else "❌ Not configured")
+    token_status += " | Write Token: " + ("✅ Available" if HF_TOKEN_WRITE else "❌ Not configured")
+    
+    token_status_class = "token-ok" if HF_TOKEN_READ else "token-warning"
+    gr.Markdown(f"""<div class="token-status {token_status_class}">{token_status}</div>""")
     
     with gr.Tab("Create Content"):
         with gr.Row():
             with gr.Column(scale=2):
-                platform = gr.Dropdown(
-                    choices=["Instagram", "TikTok", "Both"],
-                    label="Platform",
-                    value="Instagram"
-                )
+                with gr.Row():
+                    platform = gr.Dropdown(
+                        choices=["Instagram", "TikTok", "Both"],
+                        label="Platform",
+                        value="Instagram"
+                    )
+                    
+                    model_choice = gr.Dropdown(
+                        choices=["Auto-Select", "DialoGPT (Instagram)", "GPT-2 (TikTok)", "FLAN-T5"],
+                        label="AI Model",
+                        value="Auto-Select",
+                        interactive=True
+                    )
                 
                 content_type = gr.Dropdown(
                     choices=[
@@ -294,6 +484,19 @@ with gr.Blocks(
         # Output section
         output = gr.Markdown(elem_classes="output-markdown")
         
+        # Improvement UI
+        with gr.Row():
+            improve_btn = gr.Button("Improve with AI", variant="secondary")
+            enhance_tone = gr.Dropdown(
+                choices=["Casual", "Professional", "Fun", "Inspirational", "Educational", "Trendy"],
+                label="Enhance Tone",
+                value="Casual"
+            )
+            
+            # Save to Hub button (only show if write token is available)
+            if HF_TOKEN_WRITE:
+                save_btn = gr.Button("💾 Save to Hub", variant="secondary")
+        
         # Event handlers
         def update_tips(selected_platform):
             if selected_platform == "Instagram":
@@ -332,10 +535,41 @@ with gr.Blocks(
             generate_social_media_content,
             inputs=[
                 platform, content_type, topic, target_audience,
-                tone, post_length, brand_voice, key_message, call_to_action
+                tone, post_length, brand_voice, key_message, call_to_action, model_choice
             ],
             outputs=output
         )
+        
+        def improve_existing_content(content, platform, tone):
+            if content and "## Content:" in content:
+                # Extract the actual content part
+                content_part = content.split("## Content:")[1].split("## Hashtags:")[0].strip()
+                improved = improve_content(content_part, platform, tone)
+                
+                # Reconstruct the output
+                parts = content.split("## Content:")
+                parts[1] = f"\n{improved}\n\n"
+                return "## Content:".join(parts)
+            return content
+        
+        improve_btn.click(
+            improve_existing_content,
+            inputs=[output, platform, enhance_tone],
+            outputs=output
+        )
+        
+        # Save to Hub handler
+        if HF_TOKEN_WRITE:
+            def save_content_to_hub(content):
+                if not content:
+                    return "No content to save"
+                return save_to_hub(content)
+            
+            save_btn.click(
+                save_content_to_hub,
+                inputs=output,
+                outputs=gr.Markdown()
+            )
     
     with gr.Tab("Templates"):
         gr.Markdown("### Content Templates")