minor update for new models; postprocessing for md format

requirements.txt

@@ -21,4 +21,6 @@ anthropic
 openai
 cohere
 mistralai
-peft
+peft
+mdit_plain
+markdown_it

src/backend/model_operations.py

@@ -160,7 +160,7 @@ class SummaryGenerator:
         using_replicate_api = False
         replicate_api_models = ['snowflake', 'llama-3.1-405b']
         using_pipeline = False
-        pipeline_models = ['llama-3.1', 'phi-3-mini','falcon-7b', 'phi-3.5']
+        pipeline_models = ['llama-3.1', 'phi-3-mini','falcon-7b', 'phi-3.5', 'mistral-nemo']
 
         for replicate_api_model in replicate_api_models:
             if replicate_api_model in self.model_id.lower():
@@ -325,7 +325,20 @@ class SummaryGenerator:
             result = message.content[0].text
             print(result)
             return result
-                
+
+        elif 'command-r' in self.model_id.lower():
+            co = cohere.Client(os.getenv('COHERE_API_TOKEN'))
+            response = co.chat(
+                chat_history=[
+                    {"role": "SYSTEM", "message": system_prompt},
+                ],
+                message=user_prompt,
+            )
+            result = response.text
+            print(result)
+            return result
+
+
         elif 'mistral-large' in self.model_id.lower():
             api_key = os.environ["MISTRAL_API_KEY"]
             client = Mistral(api_key=api_key)
@@ -554,14 +567,8 @@ class EvaluationModel:
         for doc, summary in source_summary_pairs:
             if util.is_summary_valid(summary):
                 try:
-                    summary = summary.replace('<bos>','').replace('<eos>','').strip()
+                    summary = util.normalize_summary(summary)
                     score = self.predict([(doc, summary)])[0]
-                    # print(score)
-                    # if score < 0.5:
-                    #     print(doc)
-                    #     print('-'*10)
-                    #     print(summary)
-                    #     print('='*20)
                     hem_scores.append(score)
                     sources.append(doc)
                     summaries.append(summary)

src/backend/util.py

@@ -1,3 +1,8 @@
+import re
+from markdown_it import MarkdownIt
+from mdit_plain.renderer import RendererPlain
+
+
 def is_summary_valid(summary: str) -> bool:
     """
     Checks if the summary is valid.
@@ -76,3 +81,12 @@ def format_results(model_name: str, revision: str, precision: str,
     }
 
     return results
+
+parser = MarkdownIt(renderer_cls=RendererPlain)
+
+def normalize_summary(summary: str) -> str:
+    summary = summary.replace('<bos>','').replace('<eos>','')
+    summary = parser.render(summary)
+    summary = summary.replace('*','')
+    summary = re.sub('\s{2,}', ' ', summary)
+    return summary