Finalize + fix 3.6 tokenization and hf chat template, disable pydantic protected space warnings

Author: alpayariyak

ochat/config/__init__.py

@@ -53,7 +53,7 @@ def _v3_6_role_prefix(from_role, condition, role_start_token="", role_end_token=
                                       system_as_role=True,
                                       inference_condition="GPT4 Correct",
                                       message_prefix="\n\n"),
-        hf_chat_template="{% set loop_messages = messages %}{% for message in loop_messages %}{% if message['role'] in ['user', 'assistant'] %}{% set content = '<|start_header_id|>GPT4 Correct ' + message['role'].title() + '<|end_header_id|>\n\n' + message['content'] | trim + '<|eot_id|>' %}{% if loop.index0 == 0 %}{% set content = bos_token + content %}{% endif %}{{ content }}{% endif %}{% endfor %}{% if add_generation_prompt %}{{ '<|start_header_id|>GPT4 Correct Assistant<|end_header_id|>\n\n' }}{% endif %}",
+        hf_chat_template="{% set loop_messages = messages %}{% for message in loop_messages %}{% if message['role'] in ['user', 'assistant'] %}{% set content = '<|start_header_id|>GPT4 Correct ' + message['role'].title() + '<|end_header_id|>\n\n' + message['content'] | trim + '<|eot_id|>' %}{% elif message['role'] == 'system' %}{% set content = '<|start_header_id|>System<|end_header_id|>\n\n' + message['content'] | trim + '<|eot_id|>' %}{% endif %}{% if loop.index0 == 0 %}{% set content = bos_token + content %}{% endif %}{{ content }}{% endfor %}{% if add_generation_prompt %}{{ '<|start_header_id|>GPT4 Correct Assistant<|end_header_id|>\n\n' }}{% endif %}",
     ),
 
     # OpenChat V3.2

ochat/config/conversation_template.py

@@ -33,6 +33,7 @@ class ConversationTemplate(BaseModel):
     bos_tokens_: List[int]
     eot_tokens_: List[int]
     message_prefix_tokens_: List[int]
+    system_role_tokens_: Optional[List[int]] = []
 
     def __init__(self, **data):
         tokenizer = data["tokenizer"]
@@ -63,15 +64,14 @@ def tokenize_conversations(self, conversations: Iterable[Conversation], inferenc
             for msg in conv.items:
                 role_mappings.add((msg.role, conv.condition or default_condition))
                 all_text.append(msg.content)
-
+        
+        if self.system_as_role:
+            self.system_role_tokens_ = self.tokenizer(self.role_prefix("system", ""), add_special_tokens=False).input_ids + self.message_prefix_tokens_
+        
         sys_mappings = list(sys_mappings)
         role_mappings = list(role_mappings)
 
-        # Tokenize
-        if self.system_as_role:
-            sys_mappings = dict(zip(sys_mappings, self._tokenize([self.role_prefix(sys) for sys in sys_mappings], ignore_special=False)))
-        else:
-            sys_mappings = dict(zip(sys_mappings, self._tokenize(sys_mappings)))
+        sys_mappings = dict(zip(sys_mappings, self._tokenize(sys_mappings)))
         role_mappings = dict(zip(role_mappings, self._tokenize([self.role_prefix(*args) for args in role_mappings], ignore_special=False)))
         all_text = self._tokenize(all_text)
 
@@ -89,6 +89,9 @@ def tokenize_conversations(self, conversations: Iterable[Conversation], inferenc
 
             # System
             if conv.system:
+                tokens.extend(self.system_role_tokens_)
+                weights.extend([0.] * len(self.system_role_tokens_)) 
+                
                 system = sys_mappings[conv.system]
                 tokens.extend(system)
                 weights.extend([0.] * len(system))

ochat/config/model_config.py

@@ -1,6 +1,6 @@
 from typing import Callable, Iterable
 
-from pydantic import BaseModel
+from pydantic import BaseModel, ConfigDict
 
 
 class ModelConfig(BaseModel):
@@ -15,3 +15,5 @@ class ModelConfig(BaseModel):
     # conversation template
     conversation_template: Callable
     hf_chat_template: str = None
+
+    model_config = ConfigDict(protected_namespaces=()) # Disables warnings for the model_ namespace used abvoe