From 11028c6b4e4b5510a36575958a026676d821f441 Mon Sep 17 00:00:00 2001 From: Adam Wilson Date: Sat, 16 Aug 2025 12:33:18 -0600 Subject: [PATCH] + model support (Apple OpenELM 270M Instruct, Meta TinyLlama 1.1B Chat) --- .../apple_openelm_foundation_model.py | 33 ++++++ .../base/base_foundation_model.py | 39 +++++++ .../base/base_model_config.py | 15 +++ .../base/base_model_pipeline.py | 57 ++++++++++ .../config/apple_openelm_config.py | 10 ++ .../config/meta_tinyllama_config.py | 10 ++ .../config/microsoft_phi3mini4k_config.py | 9 ++ .../factories/foundation_model_factory.py | 56 ++++++++++ .../factory/foundation_model_factory.py | 28 +++++ .../meta_tinyllama_foundation_model.py | 36 +++++++ .../microsoft_phi3_foundation_model.py | 33 ++++++ .../pipelines/apple_openelm_pipeline.py | 16 +++ .../pipelines/meta_tinyllama_pipeline.py | 16 +++ .../pipelines/microsoft_phi3mini_pipeline.py | 15 +++ .../text_generation_foundation_model.py | 1 - src/text_generation/common/constants.py | 32 +++++- src/text_generation/common/model_id.py | 7 ++ .../dependency_injection_container.py | 2 +- .../nlp/text_generation_completion_service.py | 101 ++++++++++++++---- tests/conftest.py | 2 +- 20 files changed, 491 insertions(+), 27 deletions(-) create mode 100644 src/text_generation/adapters/foundation_models/apple_openelm_foundation_model.py create mode 100644 src/text_generation/adapters/foundation_models/base/base_foundation_model.py create mode 100644 src/text_generation/adapters/foundation_models/base/base_model_config.py create mode 100644 src/text_generation/adapters/foundation_models/base/base_model_pipeline.py create mode 100644 src/text_generation/adapters/foundation_models/config/apple_openelm_config.py create mode 100644 src/text_generation/adapters/foundation_models/config/meta_tinyllama_config.py create mode 100644 src/text_generation/adapters/foundation_models/config/microsoft_phi3mini4k_config.py create mode 100644 src/text_generation/adapters/foundation_models/factories/foundation_model_factory.py create mode 100644 src/text_generation/adapters/foundation_models/factory/foundation_model_factory.py create mode 100644 src/text_generation/adapters/foundation_models/meta_tinyllama_foundation_model.py create mode 100644 src/text_generation/adapters/foundation_models/microsoft_phi3_foundation_model.py create mode 100644 src/text_generation/adapters/foundation_models/pipelines/apple_openelm_pipeline.py create mode 100644 src/text_generation/adapters/foundation_models/pipelines/meta_tinyllama_pipeline.py create mode 100644 src/text_generation/adapters/foundation_models/pipelines/microsoft_phi3mini_pipeline.py rename src/text_generation/adapters/{ => foundation_models}/text_generation_foundation_model.py (99%) create mode 100644 src/text_generation/common/model_id.py diff --git a/src/text_generation/adapters/foundation_models/apple_openelm_foundation_model.py b/src/text_generation/adapters/foundation_models/apple_openelm_foundation_model.py new file mode 100644 index 000000000..49e8bac39 --- /dev/null +++ b/src/text_generation/adapters/foundation_models/apple_openelm_foundation_model.py @@ -0,0 +1,33 @@ +from langchain_huggingface import HuggingFacePipeline +from transformers import AutoModelForCausalLM, AutoTokenizer +from src.text_generation.adapters.foundation_models.base.base_foundation_model import BaseFoundationModel +from src.text_generation.adapters.foundation_models.config.apple_openelm_config import AppleOpenELMConfig +from src.text_generation.common.model_id import ModelId + + +class AppleOpenELMFoundationModel(BaseFoundationModel): + """Apple OpenELM 270M implementation""" + + def __init__(self, config: AppleOpenELMConfig = AppleOpenELMConfig()): + self.config = config + super().__init__() + self.MODEL_ID = ModelId.APPLE_OPENELM_270M_INSTRUCT.value + + def _load_model(self) -> None: + self.tokenizer = AutoTokenizer.from_pretrained( + self.MODEL_ID, + local_files_only=self.config.local_files_only + ) + self.model = AutoModelForCausalLM.from_pretrained( + self.MODEL_ID, + local_files_only=self.config.local_files_only + ) + + def create_pipeline(self) -> HuggingFacePipeline: + pipe = self._create_base_pipeline() + return HuggingFacePipeline( + pipeline=pipe, + pipeline_kwargs={ + "return_full_text": False, + "stop_sequence": ["", "[/INST]"] + }) \ No newline at end of file diff --git a/src/text_generation/adapters/foundation_models/base/base_foundation_model.py b/src/text_generation/adapters/foundation_models/base/base_foundation_model.py new file mode 100644 index 000000000..e1d626702 --- /dev/null +++ b/src/text_generation/adapters/foundation_models/base/base_foundation_model.py @@ -0,0 +1,39 @@ +from src.text_generation.adapters.foundation_models.base.base_model_config import BaseModelConfig +from src.text_generation.ports.abstract_foundation_model import AbstractFoundationModel + + +from transformers import pipeline + + +from abc import abstractmethod +from typing import Any + + +class BaseFoundationModel(AbstractFoundationModel): + """Base class for all foundation models""" + + def __init__(self, config: BaseModelConfig): + self.config = config + self.model = None + self.tokenizer = None + self._load_model() + + @abstractmethod + def _load_model(self) -> None: + """Load model implementation""" + pass + + def _create_base_pipeline(self) -> Any: + """Create common pipeline configuration""" + return pipeline( + "text-generation", + do_sample=True, + max_new_tokens=self.config.max_new_tokens, + model=self.model, + repetition_penalty=self.config.repetition_penalty, + temperature=self.config.temperature, + tokenizer=self.tokenizer, + use_fast=self.config.use_fast, + pad_token_id=self.tokenizer.eos_token_id, + eos_token_id=self.tokenizer.eos_token_id + ) \ No newline at end of file diff --git a/src/text_generation/adapters/foundation_models/base/base_model_config.py b/src/text_generation/adapters/foundation_models/base/base_model_config.py new file mode 100644 index 000000000..095d913ad --- /dev/null +++ b/src/text_generation/adapters/foundation_models/base/base_model_config.py @@ -0,0 +1,15 @@ +from dataclasses import dataclass +from typing import Optional + + +@dataclass +class BaseModelConfig: + """Base configuration for foundation models""" + max_new_tokens: int = 512 + temperature: float = 0.3 + repetition_penalty: float = 1.1 + use_fast: bool = True + local_files_only: bool = False + + + diff --git a/src/text_generation/adapters/foundation_models/base/base_model_pipeline.py b/src/text_generation/adapters/foundation_models/base/base_model_pipeline.py new file mode 100644 index 000000000..f1720813a --- /dev/null +++ b/src/text_generation/adapters/foundation_models/base/base_model_pipeline.py @@ -0,0 +1,57 @@ +from transformers import pipeline +from langchain.llms import HuggingFacePipeline +from typing import Dict, Any, List +from abc import ABC, abstractmethod + +class BaseModelPipeline(ABC): + def __init__(self, model, tokenizer): + self.model = model + self.tokenizer = tokenizer + + def get_common_config(self) -> Dict[str, Any]: + """Common configuration shared across all models""" + return { + "do_sample": True, + "temperature": 0.3, + "repetition_penalty": 1.1, + "use_fast": True, + "pad_token_id": self.tokenizer.eos_token_id, + "eos_token_id": self.tokenizer.eos_token_id, + } + + @abstractmethod + def get_model_specific_config(self) -> Dict[str, Any]: + """Model-specific configuration overrides""" + pass + + @abstractmethod + def get_stop_sequences(self) -> List[str]: + """Model-specific stop sequences""" + pass + + def _create_base_pipeline(self): + """Create the base pipeline with merged configurations""" + config = self.get_common_config() + config.update(self.get_model_specific_config()) + + return pipeline( + "text-generation", + model=self.model, + tokenizer=self.tokenizer, + **config + ) + + def create_pipeline(self) -> HuggingFacePipeline: + """Create the final HuggingFace pipeline""" + pipe = self._create_base_pipeline() + + return HuggingFacePipeline( + pipeline=pipe, + pipeline_kwargs={ + "return_full_text": False, + "stop_sequence": self.get_stop_sequences() + } + ) + + + diff --git a/src/text_generation/adapters/foundation_models/config/apple_openelm_config.py b/src/text_generation/adapters/foundation_models/config/apple_openelm_config.py new file mode 100644 index 000000000..ca0e866d0 --- /dev/null +++ b/src/text_generation/adapters/foundation_models/config/apple_openelm_config.py @@ -0,0 +1,10 @@ +from typing import Optional +from dataclasses import dataclass +from src.text_generation.adapters.foundation_models.base.base_model_config import BaseModelConfig + + +@dataclass +class AppleOpenELMConfig(BaseModelConfig): + """OpenELM-specific configuration""" + use_cache: bool = True + pad_token_id: Optional[int] = None \ No newline at end of file diff --git a/src/text_generation/adapters/foundation_models/config/meta_tinyllama_config.py b/src/text_generation/adapters/foundation_models/config/meta_tinyllama_config.py new file mode 100644 index 000000000..b04eccf4c --- /dev/null +++ b/src/text_generation/adapters/foundation_models/config/meta_tinyllama_config.py @@ -0,0 +1,10 @@ +from dataclasses import dataclass +from typing import Any, Dict, Optional +from src.text_generation.adapters.foundation_models.base.base_model_config import BaseModelConfig + + +@dataclass +class MetaTinyLlamaConfig(BaseModelConfig): + """TinyLlama-specific configuration""" + use_flash_attention: bool = False + rope_scaling: Optional[Dict[str, Any]] = None \ No newline at end of file diff --git a/src/text_generation/adapters/foundation_models/config/microsoft_phi3mini4k_config.py b/src/text_generation/adapters/foundation_models/config/microsoft_phi3mini4k_config.py new file mode 100644 index 000000000..70c7cd665 --- /dev/null +++ b/src/text_generation/adapters/foundation_models/config/microsoft_phi3mini4k_config.py @@ -0,0 +1,9 @@ +from dataclasses import dataclass +from src.text_generation.adapters.foundation_models.base.base_model_config import BaseModelConfig + + +@dataclass +class MicrosoftPhi3Mini4KConfig(BaseModelConfig): + """Phi3-specific configuration""" + trust_remote_code: bool = True + torch_dtype: str = "auto" \ No newline at end of file diff --git a/src/text_generation/adapters/foundation_models/factories/foundation_model_factory.py b/src/text_generation/adapters/foundation_models/factories/foundation_model_factory.py new file mode 100644 index 000000000..a7246088f --- /dev/null +++ b/src/text_generation/adapters/foundation_models/factories/foundation_model_factory.py @@ -0,0 +1,56 @@ +# Factory for creating foundation models +from typing import Optional + +from langchain_huggingface import HuggingFacePipeline +from src.text_generation.adapters.foundation_models.apple_openelm_foundation_model import AppleOpenELMFoundationModel +from src.text_generation.adapters.foundation_models.base.base_foundation_model import BaseFoundationModel +from src.text_generation.adapters.foundation_models.base.base_model_config import BaseModelConfig +from src.text_generation.adapters.foundation_models.meta_tinyllama_foundation_model import MetaTinyLlamaFoundationModel +from src.text_generation.adapters.foundation_models.microsoft_phi3_foundation_model import MicrosoftPhi3FoundationModel +from src.text_generation.adapters.foundation_models.pipelines.apple_openelm_pipeline import AppleOpenELMPipeline +from src.text_generation.adapters.foundation_models.pipelines.meta_tinyllama_pipeline import MetaTinyLlamaPipeline +from src.text_generation.adapters.foundation_models.pipelines.microsoft_phi3mini_pipeline import MicrosoftPhi3MiniPipeline +from src.text_generation.common.model_id import ModelId + + +class FoundationModelFactory: + """Factory for creating foundation model instances""" + + @staticmethod + def create_model(model_id: ModelId, config: Optional[BaseModelConfig] = None) -> BaseFoundationModel: + if config is None: + config = BaseModelConfig() + + model_map = { + ModelId.APPLE_OPENELM_270M_INSTRUCT.value: AppleOpenELMFoundationModel, + ModelId.META_TINYLLAMA_1_1B_CHAT.value: MetaTinyLlamaFoundationModel, + ModelId.MICROSOFT_PHI_3_MINI4K_INSTRUCT.value: MicrosoftPhi3FoundationModel + } + + if model_id not in model_map: + raise ValueError(f"Unsupported model type: {model_id}") + + return model_map[model_id](config) + + # Factory function to create the appropriate pipeline + def create_model_pipeline(model_id: ModelId, model, tokenizer) -> HuggingFacePipeline: + """Factory function to create the appropriate pipeline based on model name""" + + pipeline_classes = { + ModelId.APPLE_OPENELM_270M_INSTRUCT.value: AppleOpenELMPipeline, + ModelId.META_TINYLLAMA_1_1B_CHAT.value: MetaTinyLlamaPipeline, + ModelId.MICROSOFT_PHI_3_MINI4K_INSTRUCT.value: MicrosoftPhi3MiniPipeline + } + + # Determine model type from name + model_type = None + for key in pipeline_classes.keys(): + if key in model_id: + model_type = key + break + + if model_type is None: + raise ValueError(f"Unsupported model: {model_id}") + + pipeline_class = pipeline_classes[model_type] + return pipeline_class(model, tokenizer).create_pipeline() \ No newline at end of file diff --git a/src/text_generation/adapters/foundation_models/factory/foundation_model_factory.py b/src/text_generation/adapters/foundation_models/factory/foundation_model_factory.py new file mode 100644 index 000000000..a23a39449 --- /dev/null +++ b/src/text_generation/adapters/foundation_models/factory/foundation_model_factory.py @@ -0,0 +1,28 @@ +# Factory for creating foundation models +from typing import Optional +from src.text_generation.adapters.foundation_models.apple_openelm_foundation_model import AppleOpenELMFoundationModel +from src.text_generation.adapters.foundation_models.base.base_foundation_model import BaseFoundationModel +from src.text_generation.adapters.foundation_models.base.base_model_config import BaseModelConfig +from src.text_generation.adapters.foundation_models.meta_tinyllama_foundation_model import MetaTinyLlamaFoundationModel +from src.text_generation.adapters.foundation_models.microsoft_phi3_foundation_model import MicrosoftPhi3FoundationModel +from src.text_generation.common.model_id import ModelId + + +class FoundationModelFactory: + """Factory for creating foundation model instances""" + + @staticmethod + def create_model(model_id: ModelId, config: Optional[BaseModelConfig] = None) -> BaseFoundationModel: + if config is None: + config = BaseModelConfig() + + model_map = { + ModelId.APPLE_OPENELM_270M_INSTRUCT.value: AppleOpenELMFoundationModel, + ModelId.META_TINYLLAMA_1_1B_CHAT.value: MetaTinyLlamaFoundationModel, + ModelId.MICROSOFT_PHI_3_MINI4K_INSTRUCT.value: MicrosoftPhi3FoundationModel + } + + if model_id not in model_map: + raise ValueError(f"Unsupported model type: {model_id}") + + return model_map[model_id](config) \ No newline at end of file diff --git a/src/text_generation/adapters/foundation_models/meta_tinyllama_foundation_model.py b/src/text_generation/adapters/foundation_models/meta_tinyllama_foundation_model.py new file mode 100644 index 000000000..7bb47055a --- /dev/null +++ b/src/text_generation/adapters/foundation_models/meta_tinyllama_foundation_model.py @@ -0,0 +1,36 @@ +from src.text_generation.adapters.foundation_models.config.meta_tinyllama_config import MetaTinyLlamaConfig +from src.text_generation.adapters.foundation_models.base.base_foundation_model import BaseFoundationModel + + +from langchain_huggingface import HuggingFacePipeline +from transformers import AutoModelForCausalLM, AutoTokenizer + +from src.text_generation.common.model_id import ModelId + + +class MetaTinyLlamaFoundationModel(BaseFoundationModel): + """Meta TinyLlama 1.1B implementation""" + + def __init__(self, config: MetaTinyLlamaConfig = MetaTinyLlamaConfig()): + self.config = config + super().__init__() + self.MODEL_ID = ModelId.META_TINYLLAMA_1_1B_CHAT.value + + def _load_model(self) -> None: + self.tokenizer = AutoTokenizer.from_pretrained( + self.MODEL_ID, + local_files_only=self.config.local_files_only + ) + self.model = AutoModelForCausalLM.from_pretrained( + self.MODEL_ID, + local_files_only=self.config.local_files_only + ) + + def create_pipeline(self) -> HuggingFacePipeline: + pipe = self._create_base_pipeline() + return HuggingFacePipeline( + pipeline=pipe, + pipeline_kwargs={ + "return_full_text": False, + "stop_sequence": ["", "[/INST]"] + }) \ No newline at end of file diff --git a/src/text_generation/adapters/foundation_models/microsoft_phi3_foundation_model.py b/src/text_generation/adapters/foundation_models/microsoft_phi3_foundation_model.py new file mode 100644 index 000000000..ac7e78844 --- /dev/null +++ b/src/text_generation/adapters/foundation_models/microsoft_phi3_foundation_model.py @@ -0,0 +1,33 @@ +from langchain_huggingface import HuggingFacePipeline +from optimum.onnxruntime import ORTModelForCausalLM +from transformers import AutoModelForCausalLM, AutoTokenizer +from src.text_generation.adapters.foundation_models.config.microsoft_phi3mini4k_config import MicrosoftPhi3Mini4KConfig +from src.text_generation.adapters.foundation_models.base.base_foundation_model import BaseFoundationModel +from src.text_generation.common.model_id import ModelId + + +class MicrosoftPhi3FoundationModel(BaseFoundationModel): + """Microsoft Phi3 Mini 4K implementation""" + def __init__(self, config: MicrosoftPhi3Mini4KConfig = MicrosoftPhi3Mini4KConfig()): + self.config = config + super().__init__() + self.MODEL_ID = ModelId.MICROSOFT_PHI_3_MINI4K_INSTRUCT.value + + def _load_model(self) -> None: + self.tokenizer = AutoTokenizer.from_pretrained( + self.MODEL_ID, + local_files_only=self.config.local_files_only + ) + self.model = AutoModelForCausalLM.from_pretrained( + self.MODEL_ID, + local_files_only=self.config.local_files_only + ) + + def create_pipeline(self) -> HuggingFacePipeline: + pipe = self._create_base_pipeline() + return HuggingFacePipeline( + pipeline=pipe, + pipeline_kwargs={ + "return_full_text": False, + "stop_sequence": ["<|end|>", "<|user|>", ""] + }) \ No newline at end of file diff --git a/src/text_generation/adapters/foundation_models/pipelines/apple_openelm_pipeline.py b/src/text_generation/adapters/foundation_models/pipelines/apple_openelm_pipeline.py new file mode 100644 index 000000000..6952fbba2 --- /dev/null +++ b/src/text_generation/adapters/foundation_models/pipelines/apple_openelm_pipeline.py @@ -0,0 +1,16 @@ +from src.text_generation.adapters.foundation_models.base.base_model_pipeline import BaseModelPipeline + + +from typing import Any, Dict, List + + +class AppleOpenELMPipeline(BaseModelPipeline): + def get_model_specific_config(self) -> Dict[str, Any]: + return { + "max_new_tokens": 256, # Smaller model, might need fewer tokens + "temperature": 0.4, # Override common temperature for this model + "top_k": 40, # Add top-k sampling + } + + def get_stop_sequences(self) -> List[str]: + return ["", "[/INST]"] \ No newline at end of file diff --git a/src/text_generation/adapters/foundation_models/pipelines/meta_tinyllama_pipeline.py b/src/text_generation/adapters/foundation_models/pipelines/meta_tinyllama_pipeline.py new file mode 100644 index 000000000..d4392c9f3 --- /dev/null +++ b/src/text_generation/adapters/foundation_models/pipelines/meta_tinyllama_pipeline.py @@ -0,0 +1,16 @@ +from src.text_generation.adapters.foundation_models.base.base_model_pipeline import BaseModelPipeline + + +from typing import Any, Dict, List + + +class MetaTinyLlamaPipeline(BaseModelPipeline): + def get_model_specific_config(self) -> Dict[str, Any]: + return { + "max_new_tokens": 512, + # TinyLlama might need slightly different settings + "top_p": 0.9, # Add nucleus sampling for better diversity + } + + def get_stop_sequences(self) -> List[str]: + return ["", "[/INST]"] \ No newline at end of file diff --git a/src/text_generation/adapters/foundation_models/pipelines/microsoft_phi3mini_pipeline.py b/src/text_generation/adapters/foundation_models/pipelines/microsoft_phi3mini_pipeline.py new file mode 100644 index 000000000..d9fa1ab3b --- /dev/null +++ b/src/text_generation/adapters/foundation_models/pipelines/microsoft_phi3mini_pipeline.py @@ -0,0 +1,15 @@ +from src.text_generation.adapters.foundation_models.base.base_model_pipeline import BaseModelPipeline + + +from typing import Any, Dict, List + + +class MicrosoftPhi3MiniPipeline(BaseModelPipeline): + def get_model_specific_config(self) -> Dict[str, Any]: + return { + "max_new_tokens": 512, + # Remove max_length to fix the warning - max_new_tokens takes precedence + } + + def get_stop_sequences(self) -> List[str]: + return ["<|end|>", "<|user|>", ""] \ No newline at end of file diff --git a/src/text_generation/adapters/text_generation_foundation_model.py b/src/text_generation/adapters/foundation_models/text_generation_foundation_model.py similarity index 99% rename from src/text_generation/adapters/text_generation_foundation_model.py rename to src/text_generation/adapters/foundation_models/text_generation_foundation_model.py index 923fda47d..ef90c2357 100644 --- a/src/text_generation/adapters/text_generation_foundation_model.py +++ b/src/text_generation/adapters/foundation_models/text_generation_foundation_model.py @@ -6,7 +6,6 @@ from transformers import AutoTokenizer, pipeline from src.text_generation.ports.abstract_foundation_model import AbstractFoundationModel - class TextGenerationFoundationModel(AbstractFoundationModel): def __init__(self): diff --git a/src/text_generation/common/constants.py b/src/text_generation/common/constants.py index 6fe96e498..e869a88d6 100644 --- a/src/text_generation/common/constants.py +++ b/src/text_generation/common/constants.py @@ -1,15 +1,43 @@ class Constants: + # Phi-3 tokens PHI_3_ASSISTANT_START_TOKEN = "<|assistant|>" PHI_3_END_TOKEN = "<|end|>" PHI_3_SYSTEM_START_TOKEN = "<|system|>" PHI_3_USER_START_TOKEN = "<|user|>" - INPUT_VARIABLE_TOKEN = "input" + # OpenELM tokens + OPENELM_SYSTEM_START_TOKEN = "[INST]" + OPENELM_USER_START_TOKEN = "[INST]" + OPENELM_ASSISTANT_TOKEN = "[/INST]" + OPENELM_END_TOKEN = "" + + # TinyLlama tokens + TINYLLAMA_SYSTEM_START_TOKEN = "<|system|>" + TINYLLAMA_USER_START_TOKEN = "<|user|>" + TINYLLAMA_ASSISTANT_TOKEN = "<|assistant|>" + TINYLLAMA_END_TOKEN = "" + + INPUT_VARIABLE_TOKEN = "input" ALT_COMPLETION_TEXT = "I can't answer that." class PromptTemplateIds: + # Phi-3 templates PHI_3_MINI_4K_INSTRUCT__01_BASIC = "phi-3-mini-4k-instruct.01-basic" PHI_3_MINI_4K_INSTRUCT__02_ZERO_SHOT_CHAIN_OF_THOUGHT = "phi-3-mini-4k-instruct.02-zero-shot-cot" PHI_3_MINI_4K_INSTRUCT__03_FEW_SHOT_EXAMPLES = "phi-3-mini-4k-instruct.03-few-shot" PHI_3_MINI_4K_INSTRUCT__04_FEW_SHOT_RAG_PLUS_COT = "phi-3-mini-4k-instruct.04-few-shot-rag-plus-cot" - PHI_3_MINI_4K_INSTRUCT__05_REFLEXION = "phi-3-mini-4k-instruct.05-reflexion" \ No newline at end of file + PHI_3_MINI_4K_INSTRUCT__05_REFLEXION = "phi-3-mini-4k-instruct.05-reflexion" + + # OpenELM templates + OPENELM_270M_INSTRUCT__01_BASIC = "openelm-270m-instruct.01-basic" + OPENELM_270M_INSTRUCT__02_ZERO_SHOT_CHAIN_OF_THOUGHT = "openelm-270m-instruct.02-zero-shot-cot" + OPENELM_270M_INSTRUCT__03_FEW_SHOT_EXAMPLES = "openelm-270m-instruct.03-few-shot" + OPENELM_270M_INSTRUCT__04_FEW_SHOT_RAG_PLUS_COT = "openelm-270m-instruct.04-few-shot-rag-plus-cot" + OPENELM_270M_INSTRUCT__05_REFLEXION = "openelm-270m-instruct.05-reflexion" + + # TinyLlama templates + TINYLLAMA_1_1B_CHAT__01_BASIC = "tinyllama-1.1b-chat.01-basic" + TINYLLAMA_1_1B_CHAT__02_ZERO_SHOT_CHAIN_OF_THOUGHT = "tinyllama-1.1b-chat.02-zero-shot-cot" + TINYLLAMA_1_1B_CHAT__03_FEW_SHOT_EXAMPLES = "tinyllama-1.1b-chat.03-few-shot" + TINYLLAMA_1_1B_CHAT__04_FEW_SHOT_RAG_PLUS_COT = "tinyllama-1.1b-chat.04-few-shot-rag-plus-cot" + TINYLLAMA_1_1B_CHAT__05_REFLEXION = "tinyllama-1.1b-chat.05-reflexion" \ No newline at end of file diff --git a/src/text_generation/common/model_id.py b/src/text_generation/common/model_id.py new file mode 100644 index 000000000..670405f60 --- /dev/null +++ b/src/text_generation/common/model_id.py @@ -0,0 +1,7 @@ +from enum import Enum + + +class ModelId(Enum): + APPLE_OPENELM_270M_INSTRUCT = "apple/openelm-270m-instruct" + META_TINYLLAMA_1_1B_CHAT = "TinyLlama/TinyLlama-1.1B-Chat-v1.0" + MICROSOFT_PHI_3_MINI4K_INSTRUCT = "microsoft/Phi-3-mini-4k-instruct-onnx" \ No newline at end of file diff --git a/src/text_generation/dependency_injection_container.py b/src/text_generation/dependency_injection_container.py index c7b832ab9..aaab9962d 100644 --- a/src/text_generation/dependency_injection_container.py +++ b/src/text_generation/dependency_injection_container.py @@ -3,7 +3,7 @@ from dependency_injector import containers, providers from src.text_generation.adapters.embedding_model import EmbeddingModel from src.text_generation.adapters.prompt_injection_example_repository import PromptInjectionExampleRepository from src.text_generation.adapters.prompt_template_repository import PromptTemplateRepository -from src.text_generation.adapters.text_generation_foundation_model import TextGenerationFoundationModel +from src.text_generation.adapters.foundation_models.text_generation_foundation_model import TextGenerationFoundationModel from src.text_generation.entrypoints.http_api_controller import HttpApiController from src.text_generation.entrypoints.server import RestApiServer from src.text_generation.services.guidelines.abstract_security_guidelines_service import AbstractSecurityGuidelinesService diff --git a/src/text_generation/services/nlp/text_generation_completion_service.py b/src/text_generation/services/nlp/text_generation_completion_service.py index d95a6510b..e20f4a35c 100644 --- a/src/text_generation/services/nlp/text_generation_completion_service.py +++ b/src/text_generation/services/nlp/text_generation_completion_service.py @@ -1,10 +1,17 @@ +from enum import Enum +from typing import Optional, Dict, Any +import logging + from langchain.prompts import StringPromptTemplate from langchain_core.output_parsers import StrOutputParser from langchain_core.runnables import RunnablePassthrough, RunnableConfig from langchain_huggingface import HuggingFacePipeline from langchain_core.prompt_values import PromptValue +from src.text_generation.adapters.foundation_models.base.base_model_config import BaseModelConfig +from src.text_generation.adapters.foundation_models.factories.foundation_model_factory import FoundationModelFactory from src.text_generation.common.constants import Constants +from src.text_generation.common.model_id import ModelId from src.text_generation.domain.alternate_completion_result import AlternateCompletionResult from src.text_generation.domain.guidelines_result import GuidelinesResult from src.text_generation.domain.original_completion_result import OriginalCompletionResult @@ -21,11 +28,12 @@ from src.text_generation.services.utilities.abstract_llm_configuration_introspec from src.text_generation.services.utilities.abstract_response_processing_service import AbstractResponseProcessingService +logger = logging.getLogger(__name__) + class TextGenerationCompletionService( AbstractTextGenerationCompletionService): def __init__( self, - foundation_model: AbstractFoundationModel, response_processing_service: AbstractResponseProcessingService, prompt_template_service: AbstractPromptTemplateService, chain_of_thought_guidelines: AbstractSecurityGuidelinesService, @@ -34,38 +42,47 @@ class TextGenerationCompletionService( reflexion_guardrails: AbstractGeneratedTextGuardrailService, semantic_similarity_service: AbstractSemanticSimilarityService, prompt_injection_example_service: AbstractPromptInjectionExampleService, - llm_configuration_introspection_service: AbstractLLMConfigurationIntrospectionService): + llm_configuration_introspection_service: AbstractLLMConfigurationIntrospectionService, + default_model_type: ModelId = ModelId.MICROSOFT_PHI_3_MINI4K_INSTRUCT.value): + super().__init__() self.constants = Constants() - self.foundation_model_pipeline = foundation_model.create_pipeline() + + # Model management + self._current_model = None + self._current_model_id = None + self.default_model_id = default_model_type + self.factory = FoundationModelFactory() + + # Services self.response_processing_service = response_processing_service self.prompt_template_service = prompt_template_service + self.semantic_similarity_service = semantic_similarity_service + self.llm_configuration_introspection_service = llm_configuration_introspection_service - # set up semantic similarity service and supporting texts + # Set up semantic similarity service self.example_prompt_injection_completions = prompt_injection_example_service.get_all_completions() self.example_prompt_injection_prompts = prompt_injection_example_service.get_all_prompts() - self.semantic_similarity_service = semantic_similarity_service self.semantic_similarity_service.use_comparison_texts( self.example_prompt_injection_completions ) - # guidelines services - self.chain_of_thought_guidelines: AbstractSecurityGuidelinesService = chain_of_thought_guidelines - self.rag_context_guidelines: AbstractSecurityGuidelinesService = rag_context_guidelines - self.rag_plus_cot_guidelines: AbstractSecurityGuidelinesService = rag_plus_cot_guidelines + # Guidelines services + self.chain_of_thought_guidelines = chain_of_thought_guidelines + self.rag_context_guidelines = rag_context_guidelines + self.rag_plus_cot_guidelines = rag_plus_cot_guidelines - # guardrails services - self.reflexion_guardrails: AbstractGeneratedTextGuardrailService = reflexion_guardrails + # Guardrails service + self.reflexion_guardrails = reflexion_guardrails - # constants + # Constants and settings self.COSINE_SIMILARITY_RISK_THRESHOLD = 0.8 - - # default guidelines settings self._use_guidelines = False self._use_zero_shot_chain_of_thought = False self._use_rag_context = False + self._use_reflexion_guardrails = False - # dictionary dispatch for handling guidelines combinations + # Strategy map for guidelines self.guidelines_strategy_map = { (True, True): self._handle_cot_and_rag, (True, False): self._handle_cot_only, @@ -73,11 +90,43 @@ class TextGenerationCompletionService( (False, False): self._handle_without_guidelines, } - # default guardrails settings - self._use_reflexion_guardrails = False + # Load default model + self.load_model(default_model_type) - # introspection for logging - self.llm_configuration_introspection_service = llm_configuration_introspection_service + + def load_model( + self, + model_id: ModelId, + config: Optional[BaseModelConfig] = None, + force_reload: bool = False + ) -> None: + """Load a specific model""" + if (not force_reload and + self._current_model is not None and + self._current_model_id == model_id and + self._current_model.is_loaded()): + logger.info(f"Model {model_id.value} already loaded") + return + + if self._current_model is not None: + self._current_model.unload() + + self._current_model = self.factory.create_model(model_id, config) + self._current_model.load() + self._current_model_id = model_id + self.foundation_model_pipeline = self._current_model.create_pipeline() + + logger.info(f"Successfully loaded model: {model_id.value}") + + def switch_model(self, model_id: ModelId, config: Optional[BaseModelConfig] = None) -> None: + """Switch to a different model""" + self.load_model(model_id, config, force_reload=True) + + def get_current_model_info(self) -> Optional[Dict[str, Any]]: + """Get information about the currently loaded model""" + if self._current_model and self._current_model.is_loaded(): + return self._current_model.get_model_info() + return None def _process_prompt_with_guidelines_if_applicable(self, user_prompt: str): @@ -236,11 +285,19 @@ class TextGenerationCompletionService( return self._use_reflexion_guardrails - def invoke(self, user_prompt: str) -> TextGenerationCompletionResult: + def invoke(self, user_prompt: str, model_id: Optional[ModelId] = None) -> TextGenerationCompletionResult: + """Generate text using specified or current model""" if not user_prompt: raise ValueError(f"Parameter 'user_prompt' cannot be empty or None") - print(f'Using guidelines: {self.get_current_config()}') - completion_result: TextGenerationCompletionResult = self._process_prompt_with_guidelines_if_applicable(user_prompt) + + target_model_id = model_id or self._current_model_id or self.default_model_id + if (self._current_model_id != target_model_id or + self._current_model is None or + not self._current_model.is_loaded()): + self.load_model(target_model_id) + + print(f'Using model: {target_model_id.value}, guidelines: {self.get_current_config()}') + completion_result = self._process_prompt_with_guidelines_if_applicable(user_prompt) if not self._use_reflexion_guardrails: return completion_result diff --git a/tests/conftest.py b/tests/conftest.py index c66ef6fa0..45b108ce4 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -18,7 +18,7 @@ from src.text_generation import config from src.text_generation.adapters.embedding_model import EmbeddingModel from src.text_generation.adapters.prompt_injection_example_repository import PromptInjectionExampleRepository from src.text_generation.adapters.prompt_template_repository import PromptTemplateRepository -from src.text_generation.adapters.text_generation_foundation_model import TextGenerationFoundationModel +from src.text_generation.adapters.foundation_models.text_generation_foundation_model import TextGenerationFoundationModel from src.text_generation.common.constants import Constants from src.text_generation.services.guardrails.generated_text_guardrail_service import GeneratedTextGuardrailService from src.text_generation.services.guardrails.reflexion_security_guidelines_service import ReflexionSecurityGuardrailsService