ollama prototype

requirements.txt

@@ -31,4 +31,5 @@ spacy==3.7.4
 selfcheckgpt
 immutabledict
 gputil
-bitsandbytes
+bitsandbytes
+ollama

src/backend/ollama

@@ -0,0 +1,55 @@
+from typing import List, Tuple
+import torch
+import ollama
+
+from lm_eval.api.registry import register_model
+
+from src.backend.hflm_with_measurement import HFLMWithMeasurement
+
+@register_model("ollama")
+class OllamaChatTemplate(HFLMWithMeasurement):
+    def __init__(self, use_chat_template=True, **kwargs):
+        super().__init__(**kwargs)
+        self.use_chat_template = use_chat_template
+        # Initialize the ollama model and tokenizer here
+        self.model = ollama.OllamaModel.from_pretrained(kwargs['model_name_or_path'])
+        self.tokenizer = ollama.OllamaTokenizer.from_pretrained(kwargs['model_name_or_path'])
+
+    def tok_batch_encode(
+        self,
+        strings: List[str],
+        padding_side: str = "left",
+        left_truncate_len: int = None,
+        truncation: bool = False,
+    ) -> Tuple[torch.Tensor, torch.Tensor]:
+
+        if self.use_chat_template:
+            try:
+                updated_strings = []
+                for input_string in strings:
+                    messages = [
+                        {"role": "user", "content": f"{input_string}"},
+                    ]
+                    updated_string = self.tokenizer.apply_chat_template(messages, tokenize=False)
+                    updated_strings.append(updated_string)
+                strings = updated_strings[:]
+            except Exception as e:
+                print(f"Failed to update input string with chat template: {e}")
+
+        # Encode a batch of strings. Converts to tensors and pads automatically.
+        old_padding_side = self.tokenizer.padding_side
+        self.tokenizer.padding_side = padding_side
+
+        encoding = self.tokenizer(
+            strings,
+            truncation=truncation,
+            padding="longest",
+            return_tensors="pt",
+            add_special_tokens=True,
+        )
+        if left_truncate_len:
+            encoding["input_ids"] = encoding["input_ids"][:, -left_truncate_len:]
+            encoding["attention_mask"] = encoding["attention_mask"][:, -left_truncate_len:]
+        self.tokenizer.padding_side = old_padding_side
+
+        return encoding["input_ids"], encoding["attention_mask"]

src/display/utils.py

@@ -187,6 +187,7 @@ class InferenceFramework(Enum):
     # "moe-infinity", hf-chat
     MoE_Infinity = ModelDetails("moe-infinity")
     HF_Chat = ModelDetails("hf-chat")
+    Ollama = ModelDetails("ollama")
     Unknown = ModelDetails("?")
 
     def to_str(self):
@@ -198,6 +199,8 @@ class InferenceFramework(Enum):
             return InferenceFramework.MoE_Infinity
         if inference_framework in ["hf-chat"]:
             return InferenceFramework.HF_Chat
+        if inference_framework in ["ollama"]:
+            return InferenceFramework.Ollama
         return InferenceFramework.Unknown
 
 class GPUType(Enum):