Update memes.py

memes.py

@@ -29,6 +29,7 @@ def load_llama3():
     """
     Load Llama-3.2-1B and its tokenizer.
     """
+    st.write("🔄 Loading Llama-3.2-1B model and tokenizer...")
     tokenizer = AutoTokenizer.from_pretrained(
         "meta-llama/Llama-3.2-1B",
         trust_remote_code=True,
@@ -41,12 +42,16 @@ def load_llama3():
         trust_remote_code=True,
         use_auth_token=st.secrets["HUGGINGFACE_TOKEN"]
     )
+    st.write("✅ Model and tokenizer loaded.")
     return tokenizer, model
 
+
 def call_llama3(prompt: str, max_new_tokens: int = 200) -> str:
     """
     Generate text with Llama-3.2-1B.
     """
+    st.write("📝 Sending prompt to Llama-3.2-1B...")
+    st.write(prompt if len(prompt) < 200 else prompt[:200] + '...')
     tokenizer, model = load_llama3()
     inputs = tokenizer(prompt, return_tensors="pt").to(model.device)
     outputs = model.generate(
@@ -55,31 +60,52 @@ def call_llama3(prompt: str, max_new_tokens: int = 200) -> str:
         do_sample=False,
         pad_token_id=tokenizer.eos_token_id,
     )
-    return tokenizer.decode(outputs[0], skip_special_tokens=True)
+    text = tokenizer.decode(outputs[0], skip_special_tokens=True)
+    st.write("💡 Model response:")
+    st.write(text)
+    return text
+
 
 def article_to_meme(article_text: str) -> str:
-    # Summarize
+    """
+    End-to-end pipeline:
+    1) Summarize the article via Llama-3.2-1B.
+    2) Ask Llama-3.2-1B to choose a meme template and produce two 6-8 word captions.
+    3) Parse the model's response.
+    4) Call Imgflip API to render the meme and return its URL.
+    """
+    # 1) Summarize
+    st.write("▶️ Step 1: Summarizing article...")
     sum_prompt = SUMMARY_PROMPT.format(article_text=article_text)
     summary = call_llama3(sum_prompt, max_new_tokens=100).strip()
+    st.write(f"📝 Summary: {summary}")
 
-    # Meme template + captions
+    # 2) Template + captions
+    st.write("▶️ Step 2: Generating meme template and captions...")
     meme_prompt = MEME_PROMPT.format(summary=summary)
     llama_out = call_llama3(meme_prompt, max_new_tokens=150)
 
-    # Parse response
+    # 3) Parse the response
+    st.write("▶️ Step 3: Parsing model output...")
     tpl_match   = re.search(r"template:\s*(.+)", llama_out, re.IGNORECASE)
     text0_match = re.search(r"text0:\s*(.+)",    llama_out, re.IGNORECASE)
     text1_match = re.search(r"text1:\s*(.+)",    llama_out, re.IGNORECASE)
     if not (tpl_match and text0_match and text1_match):
+        st.error(f"❌ Could not parse model output:\n{llama_out}")
         raise ValueError(f"Could not parse model output:\n{llama_out}")
 
     template = tpl_match.group(1).strip()
     text0     = text0_match.group(1).strip()
     text1     = text1_match.group(1).strip()
+    st.write(f"🎯 Chosen template: {template}")
+    st.write(f"💬 text0: {text0}")
+    st.write(f"💬 text1: {text1}")
 
-    # Render meme
+    # 4) Render the meme via Imgflip
+    st.write("▶️ Step 4: Rendering meme via Imgflip...")
     template_id = TEMPLATE_IDS.get(template)
     if template_id is None:
+        st.error(f"❌ Unknown template: {template}")
         raise KeyError(f"Unknown template: {template}")
 
     creds = st.secrets["imgflip"]
@@ -91,8 +117,13 @@ def article_to_meme(article_text: str) -> str:
         "text1":       text1,
     }
     resp = requests.post(IMGFLIP_URL, params=params)
+    st.write("🔗 Imgflip API called...")
     resp.raise_for_status()
     data = resp.json()
-    if not data["success"]:
-        raise Exception(data["error_message"])
-    return data["data"]["url"]
+    if not data.get("success", False):
+        st.error(f"❌ Imgflip error: {data.get('error_message')}")
+        raise Exception(data.get("error_message"))
+
+    meme_url = data["data"]["url"]
+    st.write(f"✅ Meme URL: {meme_url}")
+    return meme_url