updated the app.py

app.py

@@ -5,88 +5,131 @@ import os
 import google.generativeai as genai
 from sentence_transformers import SentenceTransformer
 from sklearn.metrics.pairwise import cosine_similarity
+from sklearn.preprocessing import LabelEncoder
 from dotenv import load_dotenv
-
-# Import functions from model-prep
-from xgboost import XGBRegressor # Use Regressor as per model-prep
+from datasets import load_dataset
+from huggingface_hub import hf_hub_download
 import pickle
-from importlib.util import spec_from_file_location
-import sys
-# Since we are loading artifacts, we don't strictly need model-prep.py logic anymore.
-# But keeping basic imports is fine.
 
 # Load environment variables
 load_dotenv()
 
 # --- GLOBAL STATE ---
 MODEL = None
-VECTORIZER = None
 KNOWLEDGE_DF = None
 ST_MODEL = None
+ENCODERS = {} # To store label encoders
 
 def initialize_app():
-    """Initializes the model and data on app startup."""
-    global MODEL, VECTORIZER, KNOWLEDGE_DF, ST_MODEL
-    
-    print("⏳ initializing app: Loading pre-computed artifacts...")
-    
-    # 1. Load Parquet Data (Knowledge Base)
-    # We expect this file to exist now.
-    parquet_path = 'tiktok_knowledge_base.parquet'
-    if not os.path.exists(parquet_path):
-        raise FileNotFoundError(f"Required file '{parquet_path}' not found! Run model-prep.py first.")
-    
-    print(f"📂 Loading data from {parquet_path}...")
-    knowledge_df = pd.read_parquet(parquet_path)
-    
-    # 2. Load Model
-    print("🧠 Loading Prediction Model (Pickle)...")
-    with open("viral_model.pkl", "rb") as f:
-        model = pickle.load(f)
-    print(f"    -> Loaded model type: {type(model).__name__}")
-    
-    # 3. Load Vectorizer
-    print("🔤 Loading TF-IDF Vectorizer...")
-    with open("tfidf_vectorizer.pkl", "rb") as f:
-        tfidf = pickle.load(f)
-
-    # 4. Load Sentence Transformer
+    """Initializes the model and data on app startup from Hugging Face."""
+    global MODEL, KNOWLEDGE_DF, ST_MODEL, ENCODERS
+    
+    print("⏳ initializing app: Loading resources from Hugging Face Cloud...")
+    
+    # 1. Load Dataset from HF
+    print("📂 Downloading Dataset (MatanKriel/social-assitent-synthetic-data)...")
+    try:
+        dataset = load_dataset("MatanKriel/social-assitent-synthetic-data")
+        if 'train' in dataset:
+            knowledge_df = dataset['train'].to_pandas()
+        else:
+            knowledge_df = dataset.to_pandas()
+        print(f"    -> Loaded {len(knowledge_df)} rows.")
+    except Exception as e:
+        raise RuntimeError(f"Failed to load dataset from HF: {e}")
+    
+    # 2. FIT ENCODERS (Crucial for Feature Consistency)
+    print("🔤 Fitting Label Encoders...")
+    cat_cols = ['category', 'gender', 'day_of_week', 'age']
+    for c in cat_cols:
+        if c in knowledge_df.columns:
+            le = LabelEncoder()
+            # Ensure all values are strings
+            le.fit(knowledge_df[c].astype(str))
+            ENCODERS[c] = le
+            print(f"    -> Encoded '{c}': {len(le.classes_)} classes")
+        else:
+            print(f"    ⚠️ Warning: Column '{c}' missing from dataset!")
+            
+    # 3. Load Model from HF
+    print("🧠 Downloading Model (MatanKriel/social-assitent-viral-predictor)...")
+    try:
+        model_path = hf_hub_download(repo_id="MatanKriel/social-assitent-viral-predictor", filename="viral_model.pkl")
+        with open(model_path, "rb") as f:
+            model = pickle.load(f)
+        print(f"    -> Loaded model: {type(model).__name__}")
+    except Exception as e:
+        # Fallback to local
+        if os.path.exists("viral_model.pkl"):
+            print(f"    ⚠️ HF Download failed ({e}). Loading local 'viral_model.pkl' instead.")
+            with open("viral_model.pkl", "rb") as f:
+                model = pickle.load(f)
+        else:
+            raise RuntimeError(f"Failed to load model from HF and no local backup found: {e}")
+    
+    # 4. Load SentenceTransformer
     print("🔌 Loading SentenceTransformer...")
-    # device=model_prep.device might fail if we don't import model_prep executed. 
-    # Just use defaults or check pytorch standardly.
+    embedding_model_name = "sentence-transformers/all-mpnet-base-v2"
+    print(f"    -> Model: {embedding_model_name}")
+    
     import torch
     device = "mps" if torch.backends.mps.is_available() else "cpu"
-    st_model = SentenceTransformer('all-MiniLM-L6-v2', device=device)
+    st_model = SentenceTransformer(embedding_model_name, device=device)
+    
+    # 5. Generate Embeddings (On-the-fly)
+    print("⚙️ Generating Embeddings for Knowledge Base (One-time setup)...")
+    if 'embedding' not in knowledge_df.columns:
+        embeddings = st_model.encode(knowledge_df['description'].fillna("").tolist(), 
+                                   convert_to_numpy=True, 
+                                   show_progress_bar=True)
+        knowledge_df['embedding'] = list(embeddings)
+    else:
+        print("    -> Embeddings already present in dataset.")
     
     MODEL = model
-    VECTORIZER = tfidf
     KNOWLEDGE_DF = knowledge_df
     ST_MODEL = st_model
     print("✅ App initialized (Inference Mode)!")
 
-def predict_and_optimize(user_input, duration, hour, day_of_week):
+def predict_and_optimize(user_input, duration, hour, day_of_week, category, followers, age, gender):
     if not user_input:
         return "Please enter a video description.", "", "", "", ""
         
-    # --- 1. INITIAL PREDICTION ---
-    text_vec = VECTORIZER.transform([user_input]).toarray()
+    # --- 1. ENCODE INPUTS ---
+    try:
+        # Helper to encode safely
+        def safe_encode(col, val):
+            le = ENCODERS.get(col)
+            if le:
+                # If value not seen, default to first class (mostly 'Unknown' or similar if handled, else 0)
+                if val in le.classes_:
+                    return le.transform([val])[0]
+                else:
+                    return 0 # Fallback
+            return 0
+
+        cat_encoded = safe_encode('category', category)
+        gender_encoded = safe_encode('gender', gender)
+        day_encoded = safe_encode('day_of_week', day_of_week)
+        age_encoded = safe_encode('age', age)
+        
+    except Exception as e:
+        return f"Encoding Error: {str(e)}", "", "", "", ""
+
+    # --- 2. INITIAL PREDICTION ---
+    # Feature Order MUST match model-prep.py:
+    # Embeddings + [duration, hour, followers, category_enc, gender_enc, day_enc, age_enc]
+    text_vec = ST_MODEL.encode([user_input], convert_to_numpy=True)
     
-    # Process metadata inputs
-    # Map Day of Week to is_weekend
-    weekend_days = ["Saturday", "Sunday"]
-    is_weekend = 1 if day_of_week in weekend_days else 0
+    meta_vec = np.array([[duration, hour, followers, cat_encoded, gender_encoded, day_encoded, age_encoded]])
     
-    # Construct metadata vector: [duration, hour, is_weekend, hashtag_count]
-    meta_vec = np.array([[duration, hour, is_weekend, user_input.count('#')]]) 
     feat_vec = np.hstack((text_vec, meta_vec))
     
     initial_log = MODEL.predict(feat_vec)[0]
     initial_views = int(np.expm1(initial_log))
     
-    # --- 2. VECTOR SEARCH ---
-    # Filter for viral hits in knowledge base (top 25%)
+    # --- 3. VECTOR SEARCH ---
     high_perf_df = KNOWLEDGE_DF[KNOWLEDGE_DF['views'] > KNOWLEDGE_DF['views'].quantile(0.75)].copy()
-    
     user_embedding = ST_MODEL.encode([user_input], convert_to_numpy=True)
     target_embeddings = np.stack(high_perf_df['embedding'].values)
     
@@ -96,13 +139,12 @@ def predict_and_optimize(user_input, duration, hour, day_of_week):
     
     similar_videos_str = "\n\n".join([f"{i+1}. {v}" for i, v in enumerate(top_3_videos)])
 
-    # --- 3. GEMINI OPTIMIZATION ---
+    # --- 4. GEMINI OPTIMIZATION ---
     api_key = os.getenv("GEMINI_API_KEY")
     if not api_key:
         return f"{initial_views:,}", similar_videos_str, "Error: GEMINI_API_KEY not found.", "N/A", "N/A"
         
     genai.configure(api_key=api_key)
-    # Using the updated model from the user's latest change
     try: 
         llm = genai.GenerativeModel('gemini-2.5-flash-lite')
     except:
@@ -111,33 +153,31 @@ def predict_and_optimize(user_input, duration, hour, day_of_week):
     prompt = f"""
     You are a TikTok Virality Expert.
     
-    My Draft Description: "{user_input}"
-    My Metadata Plan: Duration={duration}s, Upload Time={hour}:00, Day={day_of_week}.
+    Draft: "{user_input}"
+    Niche: {category} | Creator: {age}, {gender} with {followers} followers.
+    Context: {duration}s video posted on {day_of_week} at {hour}:00.
     
-    Here are 3 successful, viral videos that are similar to my topic:
+    Viral Examples in this niche:
     1. {top_3_videos[0]}
     2. {top_3_videos[1]}
     3. {top_3_videos[2]}
     
     Task: 
-    1. Rewrite my draft description to make it go viral.
-    2. Analyze my metadata (Duration, Time, Day) and provide a specific recommendation to improve it (e.g. "Try posting at 18:00 instead" or "Make it shorter").
+    Rewrite the draft to be more viral. Add hooks and hashtags.
+    Keep it natural and relevant to the creator persona.
     
     Output Format:
     [New Description]
-    
-    [Metadata Tip]
     """
     
     try:
         response = llm.generate_content(prompt)
         improved_idea = response.text.strip()
         
-        # --- 4. RE-SCORING ---
-        # Use same metadata for the new prediction (assuming user doesn't change time/duration)
-        new_text_vec = VECTORIZER.transform([improved_idea]).toarray()
-        new_meta_vec = np.array([[duration, hour, is_weekend, improved_idea.count('#')]])
-        new_feat_vec = np.hstack((new_text_vec, new_meta_vec))
+        # --- 5. RE-SCORING ---
+        new_text_vec = ST_MODEL.encode([improved_idea], convert_to_numpy=True)
+        # Assuming metadata stays same for optimization check
+        new_feat_vec = np.hstack((new_text_vec, meta_vec))
         
         new_log = MODEL.predict(new_feat_vec)[0]
         new_views = int(np.expm1(new_log))
@@ -154,25 +194,34 @@ def predict_and_optimize(user_input, duration, hour, day_of_week):
 # --- GRADIO UI ---
 with gr.Blocks(theme=gr.themes.Soft()) as demo:
     gr.Markdown("# 🚀 Viral Content Optimizer")
-    gr.Markdown("Enter your video idea to predict its views and get AI-powered optimizations based on 2025 trends.")
+    gr.Markdown("Enter your video idea and stats to predict views and get AI-powered optimizations.")
     
     with gr.Row():
         with gr.Column(scale=1):
             input_text = gr.Textbox(
-                label="Your Video Description", 
-                placeholder="e.g., POV: trying the new grimace shake #viral",
+                label="Video Description", 
+                placeholder="POV: When you forget to mute your mic...",
                 lines=3
             )
             
             with gr.Row():
-                duration_slider = gr.Slider(minimum=5, maximum=180, value=15, step=1, label="Duration (seconds)")
-                hour_slider = gr.Slider(minimum=0, maximum=23, value=18, step=1, label="Upload Hour (24h)")
+                category = gr.Dropdown(
+                    choices=["Entertainment", "Gaming", "Fitness", "Food", "Beauty", "Tech", "Travel", "Education", "Fashion", "Health", "DIY", "Pranks"],
+                    value="Entertainment", label="Niche"
+                )
+                followers = gr.Number(value=1000, label="Follower Count", precision=0)
             
-            day_dropdown = gr.Dropdown(
-                choices=["Monday", "Tuesday", "Wednesday", "Thursday", "Friday", "Saturday", "Sunday"], 
-                value="Friday", 
-                label="Day of Week"
-            )
+            with gr.Row():
+                age = gr.Dropdown(choices=["18-24", "25-34", "35-44", "45+"], value="18-24", label="Creator Age")
+                gender = gr.Dropdown(choices=["Male", "Female"], value="Female", label="Creator Gender")
+            
+            with gr.Row():
+                duration_slider = gr.Slider(minimum=5, maximum=60, value=15, step=1, label="Duration (s)")
+                hour_slider = gr.Slider(minimum=0, maximum=23, value=18, step=1, label="Upload Hour")
+                day_dropdown = gr.Dropdown(
+                    choices=["Monday", "Tuesday", "Wednesday", "Thursday", "Friday", "Saturday", "Sunday"], 
+                    value="Friday", label="Day"
+                )
             
             with gr.Row():
                 submit_btn = gr.Button("Analyze & Optimize ⚡", variant="primary")
@@ -195,23 +244,18 @@ with gr.Blocks(theme=gr.themes.Soft()) as demo:
 
     submit_btn.click(
         fn=predict_and_optimize,
-        inputs=[input_text, duration_slider, hour_slider, day_dropdown],
+        inputs=[input_text, duration_slider, hour_slider, day_dropdown, category, followers, age, gender],
         outputs=[initial_views, similar_videos, improved_text, new_views, uplift]
     )
 
-    # Demo Button Logic: 1. Fill Form -> 2. Run Prediction
-    demo_text = "Meetings fails when You realize you forgot to turn off your mic"
-    demo_duration = 35
-    demo_hour = 19
-    demo_day = "Tuesday"
-    
+    # Demo Button Logic
     demo_btn.click(
-        fn=lambda: (demo_text, demo_duration, demo_hour, demo_day),
+        fn=lambda: ("My protein shake recipe ended up on the floor 😭 #gymfail", 15, 19, "Monday", "Fitness", 50000, "18-24", "Male"),
         inputs=None,
-        outputs=[input_text, duration_slider, hour_slider, day_dropdown]
+        outputs=[input_text, duration_slider, hour_slider, day_dropdown, category, followers, age, gender]
     ).then(
         fn=predict_and_optimize,
-        inputs=[gr.State(demo_text), gr.State(demo_duration), gr.State(demo_hour), gr.State(demo_day)],
+        inputs=[gr.State("My protein shake recipe ended up on the floor 😭 #gymfail"), gr.State(15), gr.State(19), gr.State("Monday"), gr.State("Fitness"), gr.State(50000), gr.State("18-24"), gr.State("Male")],
         outputs=[initial_views, similar_videos, improved_text, new_views, uplift]
     )