Merge pull request #21 from INTERACT-LLM/gemma

Implement Gemma (code at the time of submission for review)

Author: MinaAlmasi

configs/models.toml

@@ -10,7 +10,11 @@ name = "llama3.1:8b"
 mlx = "mlx-community/meta-Llama-3.1-8B-Instruct-4bit"
 hf = "meta-llama/Llama-3.1-8B-Instruct"
 
+[[models]]
+name = "gemma3:12b"
+hf = "google/gemma-3-12b-it"
+
 [[models]]
 name = "mistral:7b"
 mlx = "mlx-community/Mistral-7B-Instruct-v0.3-4bit"
-hf = "mistralai/Mistral-7B-Instruct-v0.3"
+hf = "mistralai/Mistral-7B-Instruct-v0.3"

pyproject.toml

@@ -21,6 +21,8 @@ dependencies = [
     "torch>=2.6.0",
     "tqdm>=4.67.1",
     "transformers>=4.48.2",
+    "pillow>=11.1.0",
+    "torchvision>=0.21.0",
 ]
 
 [build-system]
@@ -35,3 +37,6 @@ dev = [
 [tool.ruff]
 src = ["src"]
 lint.select = ["F", "E", "W", "I001"] # F for pyflakes, E + W for pycodestyle errors + warnings, #I001 for isort
+
+[tool.uv.sources]
+transformers = { git = "https://github.com/huggingface/transformers", rev = "v4.49.0-Gemma-3" }

src/interact_llm/llm/hf_gemma.py

@@ -0,0 +1,134 @@
+"""
+HF wrapper for Gemma 
+"""
+
+from pathlib import Path
+from typing import Optional
+
+import torch
+from transformers import AutoProcessor, Gemma3ForConditionalGeneration
+
+from interact_llm.data_models.chat import ChatMessage
+
+
+class ChatHFGemma:
+    """
+    Model wrapper for loading and using a HuggingFace causal language model with HF's own libraries
+    """
+
+    def __init__(
+        self,
+        model_id: str = "google/gemma-3-12b-it",
+        cache_dir: Optional[Path] = None,
+        sampling_params: Optional[dict] = None,
+        penalty_params: Optional[dict] = None,
+        max_memory: Optional[dict] = {0: "48GB", 1: "48GB"} # specify the amount of GPUs and their vrams - Gemma needs this to be able to use 2 gpus (it is too slow on a single)
+    ):
+        self.model_id = model_id
+        self.cache_dir = cache_dir
+        self.processor = None # multi-modal model does not have a tokenizer, but a processor 
+        self.model = None
+        self.sampling_params = sampling_params
+        self.penalty_params = penalty_params
+        self.max_memory = max_memory
+
+    def load(self) -> None:
+        """
+        Lazy-loading (loads model and tokenizer if not already loaded)
+        """
+        if self.processor is None:
+            self.processor = AutoProcessor.from_pretrained(
+                self.model_id, cache_dir=self.cache_dir, use_fast=True
+            )
+
+        if self.model is None:
+            try:
+                self.model = Gemma3ForConditionalGeneration.from_pretrained(
+                    self.model_id, device_map="auto", cache_dir=self.cache_dir, 
+                    max_memory=self.max_memory
+                ).eval()
+            except Exception as e:
+                print(f"Failed to load model with max_memory={self.max_memory}. Error: {e}")
+                print("Retrying without max_memory...")
+                try:
+                    self.model = Gemma3ForConditionalGeneration.from_pretrained(
+                        self.model_id, device_map="auto", cache_dir=self.cache_dir
+                    ).eval()
+                except Exception as e:
+                    print(f"Model loading failed even without max_memory. Error: {e}")
+                    raise
+
+    
+
+    def format_params(self):
+        if self.sampling_params:
+            # normalise "temp" to "temperature" (ensures you can pass temp to the model as this is how MLX/HF defines it)
+            if "temp" in self.sampling_params:
+                self.sampling_params["temperature"] = self.sampling_params.pop("temp")
+
+            kwargs = self.sampling_params
+        else:
+            kwargs = {}
+
+        if self.penalty_params:
+            kwargs.update(self.penalty_params)
+
+        return kwargs
+
+    def format_chat_for_gemma(self, chat: list[ChatMessage]) -> list[dict]:
+        formatted_chat = []
+
+        for msg in chat.messages:
+            formatted_chat.append({
+                "role": msg.role,
+                "content": [{"type": "text", "text": msg.content}]
+            })
+
+        return formatted_chat
+
+    def generate(self, chat: list[ChatMessage], max_new_tokens: int = 3000):
+        kwargs = self.format_params()
+
+        if len(kwargs) > 0:
+            do_sample = True
+        else:
+            do_sample = False
+            print(
+                "[INFO:] No sampling parameters nor penalty parameters were passed. Setting do_sample to 'False'"
+            )
+
+        self.processor.use_default_system_prompt = False # ensure no system prompt is there
+
+        formatted_chat = self.format_chat_for_gemma(chat)
+       
+        model_inputs = self.processor.apply_chat_template(
+            formatted_chat,
+            tokenize=True,
+            return_dict=True,
+            add_generation_prompt=True,
+            return_tensors = "pt",
+            # params to fix flash attn error: https://github.com/google-deepmind/gemma/issues/169
+            padding="longest",
+            pad_to_multiple_of=8, 
+        ).to(self.model.device, dtype=torch.bfloat16)
+        
+        # fix flash attn error: https://github.com/google-deepmind/gemma/issues/169
+        self.processor.tokenizer.padding_side = "left"
+
+        input_len = model_inputs["input_ids"].shape[-1]
+        
+        with torch.inference_mode(): 
+            output = self.model.generate(
+                **model_inputs,
+
+                max_new_tokens=max_new_tokens,
+                do_sample=do_sample,
+                **kwargs,
+            )
+
+        # chat (decoded output)
+        response = self.processor.decode(output[0][input_len:], skip_special_tokens=True)
+
+        chat_message = ChatMessage(role="assistant", content=response)
+
+        return chat_message

src/interact_llm/utils/model_load.py

@@ -9,6 +9,7 @@
 
 from interact_llm.llm.hf_wrapper import ChatHF
 from interact_llm.llm.mlx_wrapper import ChatMLX
+from interact_llm.llm.hf_gemma import ChatHFGemma
 
 
 def get_model_id(
@@ -72,7 +73,7 @@ def load_model_backend(
     token_path: Path = Path(__file__).parents[3] / "tokens" / "hf_token.txt",
     cache_dir: Optional[Path] = None,
     **model_kwargs,
-) -> ChatHF | ChatMLX:
+) -> ChatHF | ChatMLX | ChatHFGemma:
     """
     Loads a model based on the specified backend ("mlx" or "hf"). Will try to login to HF 
 
@@ -84,17 +85,25 @@ def load_model_backend(
             (e.g., sampling params, see documentation for ChatHF or ChatMLX)
 
     Returns:
-        ChatHF | ChatMLX: The loaded model object.
+        ChatHF | ChatMLX | ChatGemma: The loaded model object.
     """
     model_id = get_model_id(
         models_config_path=models_config_path, model_name=model_name, backend=backend
     )
 
-    # instantiate model based on backend
-    if backend == "mlx":
-        model = ChatMLX(model_id=model_id, **model_kwargs)
-    elif backend == "hf":
-        model = ChatHF(model_id=model_id, cache_dir=cache_dir, **model_kwargs)
+    if "gemma" in model_name: 
+        if backend == "mlx":
+            raise ValueError("Model is not supported in mlx yet")
+        # Gemma should be returned immediately if `backend == "hf"`
+        model = ChatHFGemma(model_id=model_id, cache_dir=cache_dir, **model_kwargs)
+    else:
+        # instantiate model based on backend
+        if backend == "mlx":
+            model = ChatMLX(model_id=model_id, **model_kwargs)
+        elif backend == "hf":
+            model = ChatHF(model_id=model_id, cache_dir=cache_dir, **model_kwargs)
+        else:
+            raise ValueError(f"Unsupported backend: {backend}")
 
     try: 
         model.load()
@@ -110,6 +119,6 @@ def load_model_backend(
             print(f"Model {model_name} loaded successfully using {backend} backend (model_id = {model_id})")
         else:
             print(f"Unexpected error occurred: {e}")
-            raise  #rReraise other unexpected errors
+            raise  # Reraise other unexpected errors
 
-    return model
+    return model

src/scripts/simulate.py

@@ -1,10 +1,5 @@
 """
 Simulate two chat LLMs talking to each other
-
-Model tested currently:
-- mlx-community/Qwen2.5-7B-Instruct-1M-4bit
-- mlx-community/meta-Llama-3.1-8B-Instruct-4bit
-
 """
 
 import argparse

src/scripts/simulate.sh

@@ -1,6 +1,6 @@
 #!/bin/bash
 
-models=("qwen2.5:7b" "llama3.1:8b" "mistral:7b")
+models=("qwen2.5:7b" "llama3.1:8b" "mistral:7b" "gemma3:12b")
 prompt_ids=("A1" "B1" "C1")
 backend="hf"