:gem: [Feature] New model enabled: gemma-7b

README.md

@@ -15,7 +15,7 @@ Project link: https://github.com/Hansimov/hf-llm-api
 ## Features
 
 - Available Models (2024/01/22): [#5](https://github.com/Hansimov/hf-llm-api/issues/5)
-  - `mistral-7b`, `mixtral-8x7b`, `nous-mixtral-8x7b`
+  - `mistral-7b`, `mixtral-8x7b`, `nous-mixtral-8x7b`, `gemma-7b`
   - Adaptive prompt templates for different models
 - Support OpenAI API format
   - Enable api endpoint via official `openai-python` package

apis/chat_api.py

@@ -54,6 +54,13 @@ class ChatAPIApp:
                     "created": 1700000000,
                     "owned_by": "NousResearch",
                 },
+                {
+                    "id": "gemma-7b",
+                    "description": "[google/gemma-7b-it]: https://huggingface.co/google/gemma-7b-it",
+                    "object": "model",
+                    "created": 1700000000,
+                    "owned_by": "Google",
+                },
             ],
         }
         return self.available_models

messagers/message_composer.py

@@ -10,6 +10,7 @@ class MessageComposer:
         "mistral-7b",
         "openchat-3.5",
         "nous-mixtral-8x7b",
+        "gemma-7b",
     ]
 
     def __init__(self, model: str = None):
@@ -19,7 +20,7 @@ class MessageComposer:
             self.model = "mixtral-8x7b"
         self.system_roles = ["system"]
         self.inst_roles = ["user", "system", "inst"]
-        self.answer_roles = ["assistant", "bot", "answer"]
+        self.answer_roles = ["assistant", "bot", "answer", "model"]
         self.default_role = "user"
 
     def concat_messages_by_role(self, messages):
@@ -63,6 +64,11 @@ class MessageComposer:
         #   Hello, who are you?<|im_end|>
         #   <|im_start|>assistant
 
+        # Google Gemma-it
+        # <start_of_turn>user
+        # How does the brain work?<end_of_turn>
+        # <start_of_turn>model
+
         self.messages = messages
         self.merged_str = ""
 
@@ -116,6 +122,29 @@ class MessageComposer:
                     )
             self.merged_str_list.append(f"GPT4 Correct Assistant:\n")
             self.merged_str = "\n".join(self.merged_str_list)
+        # https://huggingface.co/google/gemma-7b-it#chat-template
+        elif self.model in ["gemma-7b"]:
+            self.messages = self.concat_messages_by_role(messages)
+            self.merged_str_list = []
+            self.end_of_turn = "<end_of_turn>"
+            self.start_of_turn = "<start_of_turn>"
+            for message in self.messages:
+                role = message["role"]
+                content = message["content"]
+                if role in self.inst_roles:
+                    self.merged_str_list.append(
+                        f"{self.start_of_turn}user\n{content}{self.end_of_turn}"
+                    )
+                elif role in self.answer_roles:
+                    self.merged_str_list.append(
+                        f"{self.start_of_turn}model\n{content}{self.end_of_turn}"
+                    )
+                else:
+                    self.merged_str_list.append(
+                        f"{self.start_of_turn}user\n{content}{self.end_of_turn}"
+                    )
+            self.merged_str_list.append(f"{self.start_of_turn}model\n")
+            self.merged_str = "\n".join(self.merged_str_list)
         else:
             self.merged_str = "\n".join(
                 [
@@ -206,6 +235,22 @@ class MessageComposer:
             self.append_last_instruction_to_messages(
                 inst_matches_list, pair_matches_list
             )
+        # https://huggingface.co/google/gemma-7b-it#chat-template
+        elif self.model in ["gemma-7b"]:
+            pair_pattern = r"<start_of_turn>user[\s\n]*(?P<inst>[\s\S]*?)<end_of_turn>[\s\n]*<start_of_turn>model(?P<answer>[\s\S]*?)<end_of_turn>"
+            pair_matches = re.finditer(
+                pair_pattern, self.merged_str, flags=re.MULTILINE | re.IGNORECASE
+            )
+            pair_matches_list = list(pair_matches)
+            self.messages = self.convert_pair_matches_to_messages(pair_matches_list)
+            inst_pattern = r"<start_of_turn>user\n(?P<inst>[\s\S]*?)<end_of_turn>"
+            inst_matches = re.finditer(
+                inst_pattern, self.merged_str, flags=re.MULTILINE | re.IGNORECASE
+            )
+            inst_matches_list = list(inst_matches)
+            self.append_last_instruction_to_messages(
+                inst_matches_list, pair_matches_list
+            )
         else:
             self.messages = [
                 {
@@ -218,8 +263,9 @@ class MessageComposer:
 
 
 if __name__ == "__main__":
-    model = "mixtral-8x7b"
+    # model = "mixtral-8x7b"
     # model = "nous-mixtral-8x7b"
+    model = "gemma-7b"
     composer = MessageComposer(model)
     messages = [
         {

networks/message_streamer.py

@@ -12,6 +12,7 @@ class MessageStreamer:
         "mixtral-8x7b": "mistralai/Mixtral-8x7B-Instruct-v0.1",  # 72.62, fast [Recommended]
         "mistral-7b": "mistralai/Mistral-7B-Instruct-v0.2",  # 65.71, fast
         "nous-mixtral-8x7b": "NousResearch/Nous-Hermes-2-Mixtral-8x7B-DPO",
+        "gemma-7b": "google/gemma-7b-it",
         # "openchat-3.5": "openchat/openchat-3.5-1210",  # 68.89, fast
         # "zephyr-7b-beta": "HuggingFaceH4/zephyr-7b-beta",  # ❌ Too Slow
         # "llama-70b": "meta-llama/Llama-2-70b-chat-hf",  # ❌ Require Pro User
@@ -24,12 +25,14 @@ class MessageStreamer:
         "mistral-7b": "</s>",
         "nous-mixtral-8x7b": "<|im_end|>",
         "openchat-3.5": "<|end_of_turn|>",
+        "gemma-7b": "<eos>",
     }
     TOKEN_LIMIT_MAP = {
         "mixtral-8x7b": 32768,
         "mistral-7b": 32768,
         "nous-mixtral-8x7b": 32768,
         "openchat-3.5": 8192,
+        "gemma-7b": 8192,
     }
     TOKEN_RESERVED = 100