Merge pull request #56 from imoneoi/3.5_mistral

3.5 mistral

.gitignore

@@ -6,12 +6,15 @@ wandb/
 
 # Old
 old/
+temp/
+profiler/
 
 # Logs
 logs/
 
 # eval
 eval_results/
+evalplus_codegen/
 
 # All datasets
 dataset/

ochat/config/__init__.py

@@ -0,0 +1,60 @@
+from functools import partial
+
+import torch
+import transformers
+
+from ochat.config.model_config import ModelConfig
+from ochat.config.conversation_template import Message, Conversation, ConversationTemplate
+import ochat.models
+
+
+_V3_2_PREFIXES = {
+ # OpenAI mapping
+
+ "user": "User:",
+ "assistant": "Assistant:"
+}
+
+
+def _v3_2_role_prefix(from_role, condition):
+ return f"{condition} {_V3_2_PREFIXES[from_role]}".strip()
+
+
+MODEL_CONFIG_MAP = {
+ # OpenChat V3.2
+ "openchat_v3.2": ModelConfig(
+ # Model
+ model_max_context=4096,
+ model_tokenizer_create=partial(transformers.AutoTokenizer.from_pretrained,
+ use_fast=False,
+ legacy=False),
+ model_create_for_training=partial(ochat.models.LlamaForCausalLM.from_pretrained,
+ low_cpu_mem_usage=True,
+ torch_dtype=torch.bfloat16),
+
+ # Conversation Template
+ conversation_template=partial(ConversationTemplate,
+ role_prefix=_v3_2_role_prefix,
+ eot="<|end_of_turn|>",
+ inference_condition="GPT4")
+ ),
+
+ "openchat_v3.2_mistral": ModelConfig(
+ serving_aliases=("openchat_3.5", ),
+
+ # Model
+ model_max_context=8192,
+ model_tokenizer_create=partial(transformers.AutoTokenizer.from_pretrained,
+ use_fast=False,
+ legacy=True), # Mistral use legacy=True https://huggingface.co/mistralai/Mistral-7B-v0.1/blob/main/tokenizer_config.json
+ model_create_for_training=partial(ochat.models.MistralForCausalLM.from_pretrained,
+ low_cpu_mem_usage=True,
+ torch_dtype=torch.bfloat16),
+
+ # Conversation Template
+ conversation_template=partial(ConversationTemplate,
+ role_prefix=_v3_2_role_prefix,
+ eot="<|end_of_turn|>",
+ inference_condition="GPT4 Correct")
+ ),
+}

ochat/config/conversation_template.py

@@ -0,0 +1,122 @@
+from typing import Optional, Callable, Iterable, List, Dict
+
+from pydantic import BaseModel
+
+
+class Message(BaseModel):
+ role: str
+ content: str
+
+ weight: Optional[float] = None
+
+
+class Conversation(BaseModel):
+ items: List[Message]
+
+ condition: str = ""
+ system: str = ""
+
+
+class ConversationTemplate(BaseModel):
+ tokenizer: Callable
+
+ # Prompt
+ role_prefix: Callable
+ eot: str
+
+ inference_condition: Optional[str] = None
+
+ # Private
+ bos_tokens_: List[int]
+ eot_tokens_: List[int]
+
+ def __init__(self, **data):
+ tokenizer = data["tokenizer"]
+ eot = data["eot"]
+ bos_tokens_ = tokenizer("").input_ids
+ eot_tokens_ = tokenizer(eot, add_special_tokens=False).input_ids
+
+ super().__init__(**data, bos_tokens_=bos_tokens_, eot_tokens_=eot_tokens_)
+
+ def _safe_tokenize(self, strings: Iterable[str]) -> List[List[int]]:
+ return self.tokenizer(strings, split_special_tokens=True, return_attention_mask=False, add_special_tokens=False).input_ids
+
+ def tokenize_conversations(self, conversations: Iterable[Conversation], inference: bool = False, seq_level_weight: bool = False):
+ # Pre-tokenize all conversations
+ default_condition = self.inference_condition if inference else ""
+
+ sys_mappings = set()
+ role_mappings = set()
+ all_text = []
+ for conv in conversations:
+ sys_mappings.add(conv.system)
+ for msg in conv.items:
+ role_mappings.add((msg.role, conv.condition or default_condition))
+ all_text.append(msg.content)
+
+ sys_mappings = list(sys_mappings)
+ role_mappings = list(role_mappings)
+
+ # Tokenize
+ sys_mappings = dict(zip(sys_mappings, self._safe_tokenize(sys_mappings)))
+ role_mappings = dict(zip(role_mappings, self._safe_tokenize([self.role_prefix(*args) for args in role_mappings])))
+ all_text = self._safe_tokenize(all_text)
+
+ # Convert
+ result_tokens = []
+ result_weights = []
+ all_text_idx = 0
+ for conv in conversations:
+ tokens = []
+ weights = []
+
+ # bos tokens
+ tokens.extend(self.bos_tokens_)
+ weights.extend([0.] * len(self.bos_tokens_))
+
+ # System
+ if conv.system:
+ system = sys_mappings[conv.system]
+ tokens.extend(system)
+ weights.extend([0.] * len(system))
+
+ tokens.extend(self.eot_tokens_)
+ weights.extend([0.] * len(self.eot_tokens_))
+
+ # Messages
+ last_idx = len(conv.items) - 1
+ for idx, msg in enumerate(conv.items):
+ # Prefix
+ role = role_mappings[(msg.role, conv.condition or default_condition)]
+ tokens.extend(role)
+ weights.extend([0.] * len(role))
+
+ # Message
+ text = all_text[all_text_idx]
+ all_text_idx += 1
+
+ # weight
+ w = None
+ if not inference:
+ assert msg.weight is not None
+
+ w = msg.weight
+ if seq_level_weight:
+ w /= len(text) + len(self.eot_tokens_)
+
+ # Message tokens
+ tokens.extend(text)
+ weights.extend([w] * len(text))
+
+ if not (inference and idx == last_idx): # Do not add EOT on last turn during inference
+ tokens.extend(self.eot_tokens_)
+ weights.extend([w] * len(self.eot_tokens_))
+
+ # Append result
+ result_tokens.append(tokens)
+ result_weights.append(weights)
+
+ # Sanity check
+ assert all_text_idx == len(all_text)
+
+ return result_tokens, result_weights