fix: summarize example test with Gemma 2 2B.

arey/main.py

@@ -19,6 +19,7 @@
 from arey.error import AreyError
 from arey.platform.console import SignalContextManager, get_console
 from arey.play import PlayFile
+from arey.task import close
 
 
 def _generate_response(
@@ -173,6 +174,7 @@ def task(instruction: str, overrides_file: str, verbose: bool) -> int:
         lambda: (task.result and task.result.metrics),
     )
 
+    close(task)
     _print_logs(console, verbose, task.result and task.result.logs)
     return 0
 

arey/task.py

@@ -81,4 +81,9 @@ def run(task: Task, user_input: str) -> Iterator[str]:
         finish_reason=finish_reason,
         logs=stderr.getvalue(),
     )
-    model.free()
+
+
+def close(task: Task):
+    """Close a task and free the model."""
+    if model:
+        model.free()

docs/samples/summarize.yml

@@ -8,6 +8,8 @@ type: override
 tokens:
   custom:
     prompt_prefix: |
-      You are a SEO expert. Follow the instructions and provide an answer.
+      I am a SEO expert with a deep understanding of keywords and meta texts to
+      optimize a webpage for search engines. I will follow the instruction below
+      and help the user.
 
-      #Summarize and extract keywords from user message in this JSON format: {"summary": "", "keywords": []}.
+      Summarize and extract keywords from user message in this JSON format: {"summary": "", "keywords": []}.

docs/samples/summarize_file.py

@@ -1,4 +1,5 @@
 """Summarize sample script."""
+
 #!/usr/bin/env python
 import argparse
 import glob
@@ -41,8 +42,9 @@ def summarize_markdown_file(task: Task, path: str, write: bool) -> None:
         text = StringIO()
         for response in run(task, instruction):
             text.write(response)
+        formatted_text = text.getvalue().replace("```json", "").replace("```", "")
         try:
-            summary = json.loads(text.getvalue())
+            summary = json.loads(formatted_text)
             post["auto_summary"] = summary["summary"]
             post["auto_keywords"] = summary["keywords"]
             if write:
@@ -53,7 +55,7 @@ def summarize_markdown_file(task: Task, path: str, write: bool) -> None:
             console.print("---")
             console.print(f"[dim]{task.result and task.result.metrics}[/dim]")
         except Exception:
-            console.print(text.getvalue())
+            console.print(formatted_text)
             console.print_exception()