feat: expand markdown sanitization to Groq, Ollama, and HuggingFace wrappers

concordia/contrib/language_models/groq/groq_model.py

@@ -15,14 +15,15 @@
 
 """Groq language model wrapper following BaseGPTModel style."""
 
-from collections.abc import Collection, Sequence
+from collections.abc import Collection
+from collections.abc import Sequence
+import re
 from typing import Any, override
 
 from concordia.language_model import language_model
 from concordia.utils import measurements as measurements_lib
 from groq import Groq
 
-
 _MAX_MULTIPLE_CHOICE_ATTEMPTS = 20
 
 
@@ -41,6 +42,10 @@ def __init__(
     self._measurements = measurements
     self._channel = channel
 
+  def _strip_markdown(self, text_to_strip: str) -> str:
+    """Remove markdown code blocks from the text."""
+    return re.sub(r"```(?:\w+)?\n?", "", text_to_strip).strip()
+
   def _sample_text(
       self,
       prompt: str,
@@ -59,7 +64,7 @@ def _sample_text(
         {
             "role": "system",
             "content": (
-                "You always continue sentences provided "
+                "You always continue input provided "
                 "by the user and you never repeat what "
                 "the user already said."
             ),

concordia/contrib/language_models/huggingface/huggingface_model.py

@@ -18,7 +18,9 @@
   pip install gdm-concordia[huggingface]
 """
 
-from collections.abc import Collection, Sequence
+from collections.abc import Collection
+from collections.abc import Sequence
+import re
 from typing import Any, override
 
 from concordia.language_model import language_model
@@ -27,9 +29,8 @@
 from transformers import AutoModelForCausalLM
 from transformers import AutoTokenizer
 
-
 _DEFAULT_SYSTEM_MESSAGE = (
-    'You always continue sentences provided by the user and you never repeat '
+    'You always continue inputs provided by the user and you never repeat '
     'what the user already said.'
 )
 
@@ -87,6 +88,10 @@ def device(self) -> torch.device:
     """Returns the device the model is on."""
     return next(iter(self._model.parameters())).device
 
+  def _strip_markdown(self, text_to_strip: str) -> str:
+    """Remove markdown code blocks from the text."""
+    return re.sub(r'```(?:\w+)?\n?', '', text_to_strip).strip()
+
   @override
   def sample_text(
       self,
@@ -136,6 +141,7 @@ def sample_text(
     result = self._tokenizer.decode(
         outputs[0][inputs['input_ids'].shape[1] :], skip_special_tokens=True
     )
+    result = self._strip_markdown(result)
 
     for term in terminators:
       if term in result:

concordia/contrib/language_models/ollama/ollama_model.py

@@ -15,6 +15,7 @@
 """Ollama Language Model, a wrapper for models running on the local machine."""
 
 from collections.abc import Collection, Sequence
+import re
 import json
 from typing import override
 
@@ -28,13 +29,13 @@
 _DEFAULT_TEMPERATURE = 0.5
 _DEFAULT_TERMINATORS = ()
 _DEFAULT_SYSTEM_MESSAGE = (
-    "Continue the user's sentences. Never repeat their starts. For example, "
+    "Continue the user's inputs. Never repeat their starts. For example, "
     "when you see 'Bob is', you should continue the sentence after "
     "the word 'is'. Here are some more examples: 'Question: Is Jake a "
     "turtle?\nAnswer: Jake is ' should be completed as 'not a turtle.' and "
     "'Question: What is Priya doing right now?\nAnswer: Priya is currently ' "
     "should be completed as 'working on repairing the sink.'. Notice that "
-    "it is OK to be creative with how you finish the user's sentences. The "
+    "it is OK to be creative with how you finish the user's inputs. The "
     'most important thing is to always continue in the same style as the user.'
 )
 
@@ -68,6 +69,10 @@ def __init__(
     self._measurements = measurements
     self._channel = channel
 
+  def _strip_markdown(self, text_to_strip: str) -> str:
+    """Remove markdown code blocks from the text."""
+    return re.sub(r'```(?:\w+)?\n?', '', text_to_strip).strip()
+
   @override
   def sample_text(
       self,