add phi 4 mini - 4b

Author: BBC-Esq

src/constants.py

@@ -354,6 +354,17 @@
         'gated': False,
         'max_tokens': 4096,
     },
+    'Phi 4 Mini - 4b': {
+        'model': 'Phi 4 Mini - 4b',
+        'repo_id': 'microsoft/Phi-4-mini-instruct',
+        'cache_dir': 'microsoft--Phi-4-mini-instruct',
+        'cps': 222.77,
+        'vram': 4761.80,
+        'function': 'Phi4',
+        'precision': 'bfloat16',
+        'gated': False,
+        'max_new_tokens': 4096,
+    },
     'Qwen 3 - 4b': {
         'model': 'Qwen 3 - 4b',
         'repo_id': 'Qwen/Qwen3-4B',

src/module_chat.py

@@ -133,6 +133,15 @@ def get_hf_token():
     return None
 
 
+class _StopOnToken(StoppingCriteria):
+    """Stop generation when any ID in `stop_ids` is produced."""
+    def __init__(self, stop_ids):
+        self.stop_ids = set(stop_ids)
+
+    def __call__(self, input_ids, scores, **kwargs):
+        return input_ids[0, -1].item() in self.stop_ids
+
+
 class StopAfterThink(StoppingCriteria):
     def __init__(self, tokenizer):
         self.tokenizer = tokenizer
@@ -370,18 +379,6 @@ def __init__(self, generation_settings: dict, model_name: str):
 
         self.generation_settings["pad_token_id"] = self.tokenizer.eos_token_id
 
-    @torch.inference_mode()
-    def generate_response(self, inputs, remove_token_type_ids=False):
-        if remove_token_type_ids:
-            inputs.pop("token_type_ids", None)
-
-        settings = {**inputs, **self.generation_settings,
-                    "pad_token_id": self.tokenizer.eos_token_id}
-        generated = self.model.generate(**settings)
-        txt = self.tokenizer.decode(generated[0], skip_special_tokens=True)
-        txt = txt[txt.rfind("</think>") + len("</think>"):].lstrip()
-        yield txt
-
     def create_prompt(self, augmented_query: str) -> str:
         return f"""[gMASK]<sop><|system|>
 {system_message}<|user|>
@@ -423,6 +420,80 @@ def generate_response(self, inputs):
         yield from super().generate_response(inputs)
 
 
+class Phi4(BaseModel):
+    def __init__(self, generation_settings: dict, model_name: str):
+        model_info = CHAT_MODELS[model_name]
+
+        settings = copy.deepcopy(bnb_bfloat16_settings)
+        settings["model_settings"]["attn_implementation"] = "sdpa"
+        settings["model_settings"]["device_map"] = "auto"
+
+        # Pure-CPU fallback: no quant-weights on GPU, force everything to CPU
+        if not torch.cuda.is_available():
+            settings = {"tokenizer_settings": {}, "model_settings": {"device_map": "cpu"}}
+
+        super().__init__(model_info, settings, generation_settings)
+
+        self.generation_settings["pad_token_id"] = self.tokenizer.eos_token_id
+
+    def create_prompt(self, augmented_query: str) -> str:
+        return (
+            f"<|system|>{system_message}<|end|>"
+            f"<|user|>{augmented_query}<|end|><|assistant|>"
+        )
+
+    @torch.inference_mode()
+    def generate_response(self, inputs, remove_token_type_ids: bool = False):
+        if remove_token_type_ids:
+            inputs.pop("token_type_ids", None)
+
+        eos_id   = self.tokenizer.eos_token_id
+        user_id  = self.tokenizer.convert_tokens_to_ids("<|user|>")
+        assist_id = self.tokenizer.convert_tokens_to_ids("<|assistant|>")
+
+        stop_criteria = StoppingCriteriaList([_StopOnToken({user_id, eos_id})])
+
+        streamer = TextIteratorStreamer(
+            self.tokenizer,
+            skip_prompt=True,
+            skip_special_tokens=False
+        )
+
+        gen_thread = threading.Thread(
+            target=self.model.generate,
+            kwargs={**inputs,
+                    **self.generation_settings,
+                    "streamer": streamer,
+                    "eos_token_id": eos_id,
+                    "pad_token_id": eos_id,
+                    "stopping_criteria": stop_criteria},
+            daemon=True
+        )
+        gen_thread.start()
+
+        buffer, sent = "", 0
+        ASSIST, USER, END = "<|assistant|>", "<|user|>", "<|end|>"
+
+        for chunk in streamer:
+            buffer += chunk
+
+            if ASSIST in buffer:
+                buffer = buffer.split(ASSIST)[-1]
+
+            for tag in (USER, END):
+                cut = buffer.find(tag)
+                if cut != -1:
+                    buffer = buffer[:cut]
+                    streamer.break_on_eos = True
+
+            clean = buffer.replace(ASSIST, "").replace(USER, "").replace(END, "")
+
+            if len(clean) > sent:
+                yield clean[sent:]
+                sent = len(clean)
+
+        gen_thread.join()
+
 def generate_response(model_instance, augmented_query):
     prompt = model_instance.create_prompt(augmented_query)
     inputs = model_instance.create_inputs(prompt)