Update memes.py

memes.py

@@ -3,7 +3,7 @@ import streamlit as st
 import re
 import torch
 import requests
-from transformers import AutoTokenizer, AutoModelForCausalLM
+from openai import ChatCompletion
 from prompts import SUMMARY_PROMPT, MEME_PROMPT
 
 IMGFLIP_URL = "https://api.imgflip.com/caption_image"
@@ -24,106 +24,61 @@ TEMPLATE_IDS = {
     "Hide the Pain Harold":        "27813981",
 }
 
-@st.cache_resource
-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,
-        use_auth_token=st.secrets["HUGGINGFACE_TOKEN"]
-    )
-    model = AutoModelForCausalLM.from_pretrained(
-        "meta-llama/Llama-3.2-1B",
-        device_map="auto",
-        torch_dtype=torch.float16,
-        trust_remote_code=True,
-        use_auth_token=st.secrets["HUGGINGFACE_TOKEN"]
-    )
-    st.write("✅ Model and tokenizer loaded.")
-    return tokenizer, model
+# OpenAI config
+openai_api_key = st.secrets["OPENAI_API_KEY"]
 
-
-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(
-        **inputs,
-        max_new_tokens=max_new_tokens,
-        do_sample=False,
-        pad_token_id=tokenizer.eos_token_id,
+def call_openai(prompt: str) -> str:
+    """Call gpt-4o-mini once with given prompt."""
+    response = ChatCompletion.create(
+        model="gpt-4o-mini",
+        messages=[{"role": "user", "content": prompt}],
+        max_tokens=200,
+        temperature=0.7
     )
-    text = tokenizer.decode(outputs[0], skip_special_tokens=True)
-    st.write("💡 Model response:")
-    st.write(text)
-    return text
-
+    return response.choices[0].message.content.strip()
 
 def article_to_meme(article_text: str) -> str:
-    """
-    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=500).strip()
-    st.write(f"📝 Summary: {summary}")
+    st.write("⏳ Summarizing article...")
+    summary = call_openai(SUMMARY_PROMPT.format(article_text=article_text))
+    st.write("✅ Summary complete.")
 
-    # 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)
+    # 2) Choose template + captions
+    st.write("⏳ Generating meme captions...")
+    output = call_openai(MEME_PROMPT.format(summary=summary))
+    st.write("✅ Captions generated.")
 
-    # 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}")
+    # 3) Parse model output
+    match_t = re.search(r"template:\s*(.+)", output, re.IGNORECASE)
+    match0 = re.search(r"text0:\s*(.+)", output, re.IGNORECASE)
+    match1 = re.search(r"text1:\s*(.+)", output, re.IGNORECASE)
+    if not (match_t and match0 and match1):
+        raise ValueError(f"Parsing failed: {output}")
+    template = match_t.group(1).strip()
+    text0 = match0.group(1).strip()
+    text1 = match1.group(1).strip()
 
-    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}")
-
-    # 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}")
+    # 4) Render meme
+    st.write("⏳ Rendering meme...")
+    tpl_id = TEMPLATE_IDS.get(template)
+    if not tpl_id:
         raise KeyError(f"Unknown template: {template}")
-
     creds = st.secrets["imgflip"]
-    params = {
-        "template_id": template_id,
-        "username":    creds["username"],
-        "password":    creds["password"],
-        "text0":       text0,
-        "text1":       text1,
-    }
-    resp = requests.post(IMGFLIP_URL, params=params)
-    st.write("🔗 Imgflip API called...")
+    resp = requests.post(
+        IMGFLIP_URL,
+        params={
+            "template_id": tpl_id,
+            "username": creds["username"],
+            "password": creds["password"],
+            "text0": text0,
+            "text1": text1,
+        }
+    )
     resp.raise_for_status()
     data = resp.json()
     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}")
+    st.write(f"✅ Meme ready: [View here]({meme_url})")
     return meme_url