Core 1.1.0

bpm-ai-core/bpm_ai_core/classification/transformers_classifier.py

@@ -0,0 +1,63 @@
+import logging
+
+from bpm_ai_core.classification.zero_shot_classifier import ZeroShotClassifier, ClassificationResult
+
+try:
+    from transformers import pipeline, AutoTokenizer
+    has_transformers = True
+except ImportError:
+    has_transformers = False
+
+logger = logging.getLogger(__name__)
+
+DEFAULT_MODEL_EN = "MoritzLaurer/deberta-v3-large-zeroshot-v1.1-all-33"
+DEFAULT_MODEL_MULTI = "MoritzLaurer/mDeBERTa-v3-base-mnli-xnli"
+
+
+class TransformersClassifier(ZeroShotClassifier):
+    """
+    Local zero-shot classification model based on Huggingface transformers library.
+
+    To use, you should have the ``transformers`` python package installed.
+    """
+
+    def __init__(self, model: str = DEFAULT_MODEL_EN):
+        if not has_transformers:
+            raise ImportError('transformers is not installed')
+        self.model = model
+
+    def classify_with_metadata(
+            self,
+            text: str,
+            classes: list[str],
+            hypothesis_template: str | None = None
+    ) -> ClassificationResult:
+        zeroshot_classifier = pipeline("zero-shot-classification", model=self.model)
+
+        tokenizer = AutoTokenizer.from_pretrained(self.model)
+        input_tokens = len(tokenizer.encode(text))
+        max_tokens = tokenizer.model_max_length
+        logger.debug(f"Input tokens: {input_tokens}")
+        if input_tokens > max_tokens:
+            logger.warning(
+                f"Input tokens exceed max model context size: {input_tokens} > {max_tokens}. Input will be truncated."
+            )
+
+        prediction = zeroshot_classifier(
+            text,
+            classes,
+            hypothesis_template=hypothesis_template or "This example is about {}",
+            multi_label=False
+        )
+        # Zip the labels and scores together and find the label with the max score
+        labels_scores = list(zip(prediction['labels'], prediction['scores']))
+        max_label, max_score = max(labels_scores, key=lambda x: x[1])
+
+        return ClassificationResult(
+            max_label=max_label,
+            max_score=max_score,
+            labels_scores=labels_scores
+        )
+
+
+

bpm-ai-core/bpm_ai_core/classification/zero_shot_classifier.py

@@ -0,0 +1,51 @@
+from abc import ABC, abstractmethod
+from typing import Tuple
+
+from pydantic import BaseModel
+
+from bpm_ai_core.tracing.tracing import Tracing
+
+
+class ClassificationResult(BaseModel):
+    max_label: str
+    max_score: float
+    labels_scores: list[Tuple[str, float]]
+
+
+class ZeroShotClassifier(ABC):
+    """
+    Zero Shot Classification Model
+    """
+
+    @abstractmethod
+    def classify_with_metadata(
+            self,
+            text: str,
+            classes: list[str],
+            hypothesis_template: str | None = None
+    ) -> ClassificationResult:
+        pass
+
+    def classify(
+            self,
+            text: str,
+            classes: list[str],
+            confidence_threshold: float | None = None,
+            hypothesis_template: str | None = None
+    ) -> str:
+        Tracing.tracers().start_span("classification", inputs={
+            "text": text,
+            "classes": classes,
+            "confidence_threshold": confidence_threshold,
+            "hypothesis_template": hypothesis_template
+        })
+        result = self.classify_with_metadata(
+            text=text,
+            classes=classes,
+            hypothesis_template=hypothesis_template
+        )
+        Tracing.tracers().end_span(outputs={"result": result.model_dump()})
+        # Only return the label if the score is above the threshold (if given)
+        return result.max_label \
+            if not confidence_threshold or result.max_score > confidence_threshold \
+            else None

bpm-ai-core/bpm_ai_core/extractive_qa/question_answering.py

@@ -0,0 +1,47 @@
+from abc import ABC, abstractmethod
+
+from pydantic import BaseModel
+
+from bpm_ai_core.tracing.tracing import Tracing
+
+
+class QAResult(BaseModel):
+    answer: str
+    score: float
+    start_index: int
+    end_index: int
+
+
+class ExtractiveQA(ABC):
+    """
+    Extractive Question Answering Model
+    """
+
+    @abstractmethod
+    def answer_with_metadata(
+            self,
+            context: str,
+            question: str
+    ) -> QAResult:
+        pass
+
+    def answer(
+            self,
+            context: str,
+            question: str,
+            confidence_threshold: float | None = 0.1
+    ) -> str:
+        Tracing.tracers().start_span("extractive_qa", inputs={
+            "context": context,
+            "question": question,
+            "confidence_threshold": confidence_threshold
+        })
+        result = self.answer_with_metadata(
+            context=context,
+            question=question
+        )
+        Tracing.tracers().end_span(outputs={"result": result.model_dump()})
+        # Only return the answer if the score is above the threshold (if given)
+        return result.answer \
+            if not confidence_threshold or result.score > confidence_threshold \
+            else None

bpm-ai-core/bpm_ai_core/extractive_qa/transformers_qa.py

@@ -0,0 +1,51 @@
+import logging
+
+from bpm_ai_core.extractive_qa.question_answering import ExtractiveQA, QAResult
+
+try:
+    from transformers import pipeline, AutoTokenizer
+    has_transformers = True
+except ImportError:
+    has_transformers = False
+
+logger = logging.getLogger(__name__)
+
+
+class TransformersExtractiveQA(ExtractiveQA):
+    """
+    Local extractive question answering model based on Huggingface transformers library.
+
+    To use, you should have the ``transformers`` python package installed.
+    """
+
+    def __init__(self, model: str = "deepset/deberta-v3-large-squad2"):
+        if not has_transformers:
+            raise ImportError('transformers is not installed')
+        self.model = model
+
+    def answer_with_metadata(
+            self,
+            context: str,
+            question: str
+    ) -> QAResult:
+        qa_model = pipeline("question-answering", model=self.model)
+
+        tokenizer = AutoTokenizer.from_pretrained(self.model)
+        tokens = tokenizer.encode(context + question)
+        logger.debug(f"Input tokens: {len(tokens)}")
+
+        prediction = qa_model(
+            question=question,
+            context=context
+        )
+        logger.debug(f"prediction: {prediction}")
+
+        return QAResult(
+            answer=prediction['answer'],
+            score=prediction['score'],
+            start_index=prediction['start'],
+            end_index=prediction['end'],
+        )
+
+
+

bpm-ai-core/bpm_ai_core/llm/common/llm.py

@@ -26,7 +26,7 @@ def __init__(
         self.max_retries = max_retries
         self.retryable_exceptions = retryable_exceptions or [Exception]
 
-    def predict(
+    async def predict(
         self,
         prompt: Prompt,
         output_schema: dict[str, Any] | None = None,
@@ -44,13 +44,13 @@ def predict(
         ):
             with attempt:
                 Tracing.tracers().start_llm_trace(self, messages, attempt.retry_state.attempt_number, tools)
-                completion = self._predict(messages, output_schema, tools)
+                completion = await self._predict(messages, output_schema, tools)
                 Tracing.tracers().end_llm_trace(completion)
 
         return completion
 
     @abstractmethod
-    def _predict(
+    async def _predict(
         self,
         messages: list[ChatMessage],
         output_schema: dict[str, Any] | None = None,

bpm-ai-core/bpm_ai_core/llm/openai_chat.py

@@ -7,9 +7,21 @@
 from bpm_ai_core.util.openai import messages_to_openai_dicts, json_schema_to_openai_function
 
 try:
-    from openai import OpenAI, APIConnectionError, InternalServerError, RateLimitError
-    from openai.types.chat import ChatCompletionMessage
+    from openai import AsyncOpenAI, APIConnectionError, InternalServerError, RateLimitError
+    from openai.types.chat import ChatCompletionMessage, ChatCompletion
+    import httpx
+
     has_openai = True
+
+    client = AsyncOpenAI(
+        http_client=httpx.AsyncClient(
+            limits=httpx.Limits(
+                max_connections=1000,
+                max_keepalive_connections=100
+            )
+        ),
+        max_retries=0  # we use own retry logic
+    )
 except ImportError:
     has_openai = False
 
@@ -24,11 +36,10 @@ class ChatOpenAI(LLM):
 
     def __init__(
         self,
-        model: str = "gpt-3.5-turbo-1106",
+        model: str = "gpt-3.5-turbo-0125",
         temperature: float = 0.0,
         seed: Optional[int] = None,
-        max_retries: int = 8,
-        client_kwargs: Optional[Dict[str, Any]] = None
+        max_retries: int = 8
     ):
         if not has_openai:
             raise ImportError('openai is not installed')
@@ -41,12 +52,8 @@ def __init__(
             ]
         )
         self.seed = seed
-        self.client = OpenAI(
-            max_retries=0,  # we use own retry logic
-            **(client_kwargs or {})
-        )
 
-    def _predict(
+    async def _predict(
         self,
         messages: List[ChatMessage],
         output_schema: Optional[Dict[str, Any]] = None,
@@ -57,7 +64,7 @@ def _predict(
             tools = [Tool.from_callable(name="store_result", description="Stores your result", args_schema=output_schema)]
         if tools:
             openai_tools = [json_schema_to_openai_function(f.name, f.description, f.args_schema) for f in tools]
-        completion = self._run_completion(messages, openai_tools)
+        completion = await self._run_completion(messages, openai_tools)
 
         message = completion.choices[0].message
         if message.tool_calls:
@@ -68,7 +75,7 @@ def _predict(
         else:
             return ChatMessage(role=message.role, content=message.content)
 
-    def _run_completion(self, messages: List[ChatMessage], functions: List[dict]):
+    async def _run_completion(self, messages: List[ChatMessage], functions: List[dict]) -> ChatCompletion:
         args = {
             "model": self.model,
             "temperature": self.temperature,
@@ -82,7 +89,7 @@ def _run_completion(self, messages: List[ChatMessage], functions: List[dict]):
                    "tools": functions
                } if functions else {})
         }
-        return self.client.chat.completions.create(**args)
+        return await client.chat.completions.create(**args)
 
     @staticmethod
     def _openai_tool_calls_to_tool_message(message: ChatCompletionMessage, tools: List[Tool]) -> ToolCallsMessage:

bpm-ai-core/bpm_ai_core/pos/pos_tagger.py

@@ -0,0 +1,16 @@
+from abc import ABC, abstractmethod
+from typing import Tuple
+
+
+class POSTagger(ABC):
+    """
+    Part-of-Speech Tagging Model
+    """
+
+    @abstractmethod
+    def tag(self, text: str) -> list[Tuple[str, str]]:
+        """
+        Returns a list of tuples (token, tag). Example:
+        [('I', 'PRON'), ('am', 'AUX'), ('30', 'NUM'), ('years', 'NOUN'), ('old', 'ADJ'), ('.', 'PUNCT')]
+        """
+        pass