Update memes.py

memes.py

@@ -25,27 +25,27 @@ TEMPLATE_IDS = {
 }
 
 @st.cache_resource
-def load_mpt():
+def load_llama3():
     """
-    Load the MosaicML MPT-7B-Chat model and tokenizer.
+    Load Llama-3.2-1B and its tokenizer.
     """
     tokenizer = AutoTokenizer.from_pretrained(
-        "mosaicml/mpt-7b-chat",
+        "meta-llama/Llama-3.2-1B",
         trust_remote_code=True
     )
     model = AutoModelForCausalLM.from_pretrained(
-        "mosaicml/mpt-7b-chat",
+        "meta-llama/Llama-3.2-1B",
         device_map="auto",
         torch_dtype=torch.float16,
         trust_remote_code=True
     )
     return tokenizer, model
 
-def call_mpt(prompt: str, max_new_tokens: int = 200) -> str:
+def call_llama3(prompt: str, max_new_tokens: int = 200) -> str:
     """
-    Generate text from MPT-7B-Chat given a prompt.
+    Generate text with Llama-3.2-1B.
     """
-    tokenizer, model = load_mpt()
+    tokenizer, model = load_llama3()
     inputs = tokenizer(prompt, return_tensors="pt").to(model.device)
     outputs = model.generate(
         **inputs,
@@ -56,33 +56,26 @@ def call_mpt(prompt: str, max_new_tokens: int = 200) -> str:
     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
+    # Summarize
     sum_prompt = SUMMARY_PROMPT.format(article_text=article_text)
-    summary = call_mpt(sum_prompt, max_new_tokens=100).strip()
+    summary = call_llama3(sum_prompt, max_new_tokens=100).strip()
 
-    # 2) Template + captions
+    # Meme template + captions
     meme_prompt = MEME_PROMPT.format(summary=summary)
-    mpt_out = call_mpt(meme_prompt, max_new_tokens=150)
+    llama_out = call_llama3(meme_prompt, max_new_tokens=150)
 
-    # 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)
+    # Parse response
+    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):
-        raise ValueError(f"Could not parse model output:\n{mpt_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()
 
-    # 4) Render the meme via Imgflip
+    # Render meme
     template_id = TEMPLATE_IDS.get(template)
     if template_id is None:
         raise KeyError(f"Unknown template: {template}")
@@ -100,5 +93,4 @@ def article_to_meme(article_text: str) -> str:
     data = resp.json()
     if not data["success"]:
         raise Exception(data["error_message"])
-
     return data["data"]["url"]