Update memes.py

memes.py

@@ -25,59 +25,73 @@ TEMPLATE_IDS = {
 }
 
 @st.cache_resource
-def load_gemma():
+def load_mpt():
+    """
+    Load the MosaicML MPT-7B-Chat model and tokenizer.
+    """
     tokenizer = AutoTokenizer.from_pretrained(
-        "google/gemma-3-27b-it",
-        use_auth_token=st.secrets["HUGGINGFACE_TOKEN"]
+        "mosaicml/mpt-7b-chat",
+        trust_remote_code=True
     )
     model = AutoModelForCausalLM.from_pretrained(
-        "google/gemma-3-27b-it",
-        torch_dtype=torch.float16,
+        "mosaicml/mpt-7b-chat",
         device_map="auto",
-        use_auth_token=st.secrets["HUGGINGFACE_TOKEN"]
+        torch_dtype=torch.float16,
+        trust_remote_code=True
     )
     return tokenizer, model
 
-def call_gemma(prompt: str, max_new_tokens=200) -> str:
-    tok, model = load_gemma()
-    inputs = tok(prompt, return_tensors="pt").to(model.device)
-    out = model.generate(
+def call_mpt(prompt: str, max_new_tokens: int = 200) -> str:
+    """
+    Generate text from MPT-7B-Chat given a prompt.
+    """
+    tokenizer, model = load_mpt()
+    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=tok.eos_token_id
+        pad_token_id=tokenizer.eos_token_id,
     )
-    return tok.decode(out[0], skip_special_tokens=True)
+    return tokenizer.decode(outputs[0], skip_special_tokens=True)
 
 def article_to_meme(article_text: str) -> str:
+    """
+    End-to-end pipeline:
+    1) Summarize the article via MPT-7B-Chat.
+    2) Ask MPT-7B-Chat 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
     sum_prompt = SUMMARY_PROMPT.format(article_text=article_text)
-    summary = call_gemma(sum_prompt, max_new_tokens=100).strip()
+    summary = call_mpt(sum_prompt, max_new_tokens=100).strip()
 
-    # 2) Meme‐text generation + template choice
+    # 2) Template + captions
     meme_prompt = MEME_PROMPT.format(summary=summary)
-    gemma_out = call_gemma(meme_prompt, max_new_tokens=150)
+    mpt_out = call_mpt(meme_prompt, max_new_tokens=150)
 
-    # 3) Parse Gemini’s response
-    tpl_match   = re.search(r"template:\s*(.+)", gemma_out, re.IGNORECASE)
-    text0_match = re.search(r"text0:\s*(.+)",    gemma_out, re.IGNORECASE)
-    text1_match = re.search(r"text1:\s*(.+)",    gemma_out, re.IGNORECASE)
+    # 3) Parse the response
+    tpl_match   = re.search(r"template:\s*(.+)", mpt_out, re.IGNORECASE)
+    text0_match = re.search(r"text0:\s*(.+)",    mpt_out, re.IGNORECASE)
+    text1_match = re.search(r"text1:\s*(.+)",    mpt_out, re.IGNORECASE)
     if not (tpl_match and text0_match and text1_match):
-        raise ValueError(f"Could not parse Gemini output:\n{gemma_out}")
+        raise ValueError(f"Could not parse model output:\n{mpt_out}")
 
     template = tpl_match.group(1).strip()
     text0     = text0_match.group(1).strip()
     text1     = text1_match.group(1).strip()
 
-    # 4) Render meme
+    # 4) Render the meme via Imgflip
     template_id = TEMPLATE_IDS.get(template)
     if template_id is None:
         raise KeyError(f"Unknown template: {template}")
 
+    creds = st.secrets["imgflip"]
     params = {
         "template_id": template_id,
-        "username":    st.secrets["IMGFLIP_USERNAME"],
-        "password":    st.secrets["IMGFLIP_PASSWORD"],
+        "username":    creds["username"],
+        "password":    creds["password"],
         "text0":       text0,
         "text1":       text1,
     }
@@ -86,4 +100,5 @@ def article_to_meme(article_text: str) -> str:
     data = resp.json()
     if not data["success"]:
         raise Exception(data["error_message"])
+
     return data["data"]["url"]