crfm-helm 0.5.3__py3-none-any.whl → 0.5.5__py3-none-any.whl

helm/clients/ibm_client.py

@@ -0,0 +1,267 @@
+from abc import ABC
+from abc import abstractmethod
+
+from helm.common.hierarchical_logger import hlog
+from helm.common.cache import CacheConfig
+from helm.common.request import (
+    Request,
+    RequestResult,
+    Token,
+    wrap_request_time,
+    EMBEDDING_UNAVAILABLE_REQUEST_RESULT,
+    GeneratedOutput,
+)
+
+from helm.clients.client import CachingClient
+from helm.common.optional_dependencies import handle_module_not_found_error
+from typing import TypeVar, Generic
+from typing import Any, Dict, List
+from threading import Semaphore, Lock
+import threading
+
+try:
+    from ibm_watsonx_ai import Credentials
+    from ibm_watsonx_ai.foundation_models import ModelInference
+    from ibm_watsonx_ai.metanames import GenTextParamsMetaNames as GenParams
+    from ibm_watsonx_ai.foundation_models.schema import (
+        TextChatParameters,
+        TextGenParameters,
+        ReturnOptionProperties,
+    )
+
+except ModuleNotFoundError as e:
+    handle_module_not_found_error(e, ["ibm"])
+
+# Define the maximum number of parallel executions is limited by IBM API
+MAX_CONCURRENT_REQUESTS = 8
+__semaphores: Dict[str, Semaphore] = dict()
+__semaphores_lock = Lock()
+
+
+def _get_semaphore(model: str) -> Semaphore:
+    with __semaphores_lock:
+        if model not in __semaphores:
+            __semaphores[model] = threading.Semaphore(MAX_CONCURRENT_REQUESTS)
+
+    return __semaphores[model]
+
+
+T = TypeVar("T", TextGenParameters, TextChatParameters)
+
+
+class ModelInferenceHandler(ABC, Generic[T]):
+    @abstractmethod
+    def __init__(self, inference_engine: ModelInference):
+        """
+        :type inference_engine: object
+        """
+        self.inference_engine = inference_engine
+
+    @abstractmethod
+    def serve_request(self, prompt: str, params: T) -> Dict:
+        pass
+
+    @abstractmethod
+    def parse_response(self, response: dict) -> List[GeneratedOutput]:
+        pass
+
+    @abstractmethod
+    def create_params(self, request: Request) -> T:
+        pass
+
+
+class GenerateInferenceHandler(ModelInferenceHandler[TextGenParameters]):
+
+    def __init__(self, inference_engine: ModelInference):
+        self.inference_engine = inference_engine
+
+    def create_params(self, request: Request) -> TextGenParameters:
+        def set_temperature_requirements():
+            # Default temperature 0.05 required by ibm/granite-13b-instruct-v2
+            if self.inference_engine.model_id == "ibm/granite-13b-instruct-v2":
+                return 0.05
+            return 1e-7 if request.temperature == 0 else request.temperature
+
+        return TextGenParameters(
+            temperature=set_temperature_requirements(),
+            top_p=request.top_p,
+            max_new_tokens=request.max_tokens,
+            return_options=ReturnOptionProperties(
+                input_text=True,
+                generated_tokens=True,
+                input_tokens=False,
+                token_logprobs=True,
+                token_ranks=False,
+            ),
+            include_stop_sequence=False,
+            prompt_variables=None,
+        )
+
+    def serve_request(self, prompt: str, params: TextGenParameters) -> Dict:
+        semaphore = _get_semaphore(self.inference_engine.model_id)
+
+        with semaphore:
+            response = self.inference_engine.generate(
+                prompt=prompt,
+                params=params,
+            )
+        return response
+
+    def parse_response(self, response: dict) -> List[GeneratedOutput]:
+        completions = []
+        try:
+            for r in response["results"]:
+                sequence_logprob: float = 0
+                tokens: List[Token] = []
+                generated_text = r["generated_text"]
+                for token_and_logprob in r["generated_tokens"]:
+                    logprob = token_and_logprob.get("logprob", 0)
+                    text = token_and_logprob["text"]
+                    tokens.append(Token(text=text, logprob=logprob))
+                    sequence_logprob += logprob
+
+                completion = GeneratedOutput(text=generated_text, logprob=sequence_logprob, tokens=tokens)
+                completions.append(completion)
+        except Exception as e:
+            hlog(f"GenerateInferenceHandler failed with exception {e} during parse_response {response}")
+        return completions
+
+
+class ChatModelInferenceHandler(ModelInferenceHandler[TextChatParameters]):
+    def __init__(self, inference_engine: ModelInference):
+        self.inference_engine = inference_engine
+
+    def create_params(self, request: Request) -> TextChatParameters:
+        return TextChatParameters(
+            logprobs=True,
+            presence_penalty=0,
+            frequency_penalty=0,
+            temperature=request.temperature,
+            max_tokens=request.max_tokens,
+            top_p=request.top_p,
+        )
+
+    def parse_response(self, response: dict) -> List[GeneratedOutput]:
+        completions = []
+        try:
+            for raw_completion in response["choices"]:
+                sequence_logprob: float = 0
+                tokens: List[Token] = []
+                generated_text = raw_completion["message"]["content"]
+
+                for token_and_logprob in raw_completion["logprobs"]["content"]:
+                    logprob = token_and_logprob["logprob"]
+                    text = token_and_logprob["token"]
+                    tokens.append(Token(text=text, logprob=logprob))
+                    sequence_logprob += logprob
+
+                completion = GeneratedOutput(text=generated_text, logprob=sequence_logprob, tokens=tokens)
+                completions.append(completion)
+        except Exception as e:
+            hlog(f"ChatModelInferenceHandler failed with exception {e} during parse_response {response}")
+        return completions
+
+    def serve_request(self, prompt: str, params: TextChatParameters) -> Dict:
+        semaphore = _get_semaphore(self.inference_engine.model_id)
+
+        with semaphore:
+            response = self.inference_engine.chat(
+                messages=[{"role": "user", "content": prompt}],
+                params=params,
+            )
+        return response
+
+
+class IbmClient(CachingClient, ABC):
+    def __init__(
+        self,
+        cache_config: CacheConfig,
+        api_key: str,
+        region: str,
+        location: dict,
+        watsonx_model_name: str,
+        **kwargs,
+    ):
+        super().__init__(cache_config=cache_config)
+        self.project_id = None
+        self.url = None
+        self.watsonx_model_name = watsonx_model_name
+        self.api_key = api_key
+        self.region = region
+        self.kwargs = kwargs
+        for entry in location:
+            if entry["region"].lower() == self.region.lower():
+                self.project_id = entry["project_id"]
+                self.url = entry["url"]
+
+        assert self.project_id is not None, (
+            "Missed project_id for specified region configuration in credentials.conf, should be in list "
+            "of JSON objects with 'region', 'url', 'project_id' per region"
+        )
+        assert self.url is not None, (
+            "Missed url for specified region configuration in credentials.conf, should be in list "
+            "of JSON objects with 'region', 'url', 'project_id' per region"
+        )
+
+        self.inference_engine = ModelInference(
+            model_id=self.watsonx_model_name,
+            params={GenParams.MAX_NEW_TOKENS: 2000},
+            credentials=Credentials(api_key=api_key, url=self.url),
+            project_id=self.project_id,
+        )
+
+        hlog("Started IBM Client")
+
+    @abstractmethod
+    def make_request(self, request: Request) -> RequestResult:
+        pass
+
+    def do_call(self, inference_handler: ModelInferenceHandler, request: Request) -> RequestResult:
+        params = inference_handler.create_params(request=request)
+
+        def do_it() -> Dict[str, Any]:
+            return inference_handler.serve_request(prompt=request.prompt, params=params)
+
+        raw_request = {"prompt": request.prompt, "params": params.to_dict(), "model": request.model}
+
+        cache_key = CachingClient.make_cache_key(raw_request, request)
+        response, cached = self.cache.get(cache_key, wrap_request_time(do_it))
+        completions = inference_handler.parse_response(response)
+        return RequestResult(
+            success=True,
+            cached=cached,
+            request_time=response["request_time"],
+            request_datetime=response.get("request_datetime"),
+            completions=completions,
+            embedding=[],
+        )
+
+
+class IbmChatClient(IbmClient):
+
+    def make_request(self, request: Request) -> RequestResult:
+        # Embedding not supported for this model
+        if request.embedding:
+            return EMBEDDING_UNAVAILABLE_REQUEST_RESULT
+        try:
+            return self.do_call(
+                inference_handler=ChatModelInferenceHandler(inference_engine=self.inference_engine), request=request
+            )
+
+        except Exception as e:
+            error: str = f"IBM Chat client Model error: {e}"
+            return RequestResult(success=False, cached=False, error=error, completions=[], embedding=[])
+
+
+class IbmTextClient(IbmClient):
+    def make_request(self, request: Request) -> RequestResult:
+        # Embedding not supported for this model
+        if request.embedding:
+            return EMBEDDING_UNAVAILABLE_REQUEST_RESULT
+        try:
+            return self.do_call(
+                inference_handler=GenerateInferenceHandler(inference_engine=self.inference_engine), request=request
+            )
+        except Exception as e:
+            error: str = f"IBM Text client Model error: {e}"
+            return RequestResult(success=False, cached=False, error=error, completions=[], embedding=[])

helm/clients/image_generation/adobe_vision_client.py

@@ -9,7 +9,7 @@ from helm.common.tokenization_request import (
     DecodeRequestResult,
 )
 from helm.clients.client import Client, CachingClient
-from .image_generation_client_utils import get_single_image_multimedia_object
+from helm.clients.image_generation.image_generation_client_utils import get_single_image_multimedia_object
 
 
 class AdobeVisionClient(Client):

helm/clients/image_generation/aleph_alpha_image_generation_client.py

@@ -9,7 +9,7 @@ from helm.common.tokenization_request import (
     DecodeRequestResult,
 )
 from helm.clients.client import Client, CachingClient
-from .image_generation_client_utils import get_single_image_multimedia_object
+from helm.clients.image_generation.image_generation_client_utils import get_single_image_multimedia_object
 
 
 class AlephAlphaImageGenerationClient(Client):

helm/clients/image_generation/cogview2/sr_pipeline/__init__.py

@@ -6,9 +6,9 @@
 @Contact :   dm18@mails.tsinghua.edu.cn
 """
 
-from .direct_sr import DirectSuperResolution
-from .iterative_sr import IterativeSuperResolution
-from .sr_group import SRGroup
+from helm.clients.image_generation.cogview2.sr_pipeline.direct_sr import DirectSuperResolution
+from helm.clients.image_generation.cogview2.sr_pipeline.iterative_sr import IterativeSuperResolution
+from helm.clients.image_generation.cogview2.sr_pipeline.sr_group import SRGroup
 
 DirectSuperResolution
 IterativeSuperResolution

helm/clients/image_generation/cogview2/sr_pipeline/direct_sr.py

@@ -10,8 +10,11 @@
 import torch
 from icetk import icetk as tokenizer
 
-from .dsr_sampling import filling_sequence_dsr, IterativeEntfilterStrategy
-from .dsr_model import DsrModel
+from helm.clients.image_generation.cogview2.sr_pipeline.dsr_sampling import (
+    filling_sequence_dsr,
+    IterativeEntfilterStrategy,
+)
+from helm.clients.image_generation.cogview2.sr_pipeline.dsr_model import DsrModel
 from helm.common.optional_dependencies import handle_module_not_found_error
 
 

helm/clients/image_generation/cogview2/sr_pipeline/iterative_sr.py

@@ -8,8 +8,11 @@
 import torch
 from icetk import icetk as tokenizer
 
-from .itersr_sampling import filling_sequence_itersr, IterativeEntfilterStrategy
-from .itersr_model import ItersrModel
+from helm.clients.image_generation.cogview2.sr_pipeline.itersr_sampling import (
+    filling_sequence_itersr,
+    IterativeEntfilterStrategy,
+)
+from helm.clients.image_generation.cogview2.sr_pipeline.itersr_model import ItersrModel
 from helm.common.optional_dependencies import handle_module_not_found_error
 
 

helm/clients/image_generation/cogview2/sr_pipeline/sr_group.py

@@ -7,8 +7,8 @@
 """
 
 # here put the import lib
-from .direct_sr import DirectSuperResolution
-from .iterative_sr import IterativeSuperResolution
+from helm.clients.image_generation.cogview2.sr_pipeline.direct_sr import DirectSuperResolution
+from helm.clients.image_generation.cogview2.sr_pipeline.iterative_sr import IterativeSuperResolution
 
 from helm.common.optional_dependencies import handle_module_not_found_error
 

helm/clients/image_generation/cogview2_client.py

@@ -20,7 +20,7 @@ from helm.common.tokenization_request import (
 )
 from helm.clients.client import Client, CachingClient
 from helm.clients.image_generation.cogview2.coglm_strategy import CoglmStrategy
-from .image_generation_client_utils import get_single_image_multimedia_object
+from helm.clients.image_generation.image_generation_client_utils import get_single_image_multimedia_object
 
 
 class CogView2Client(Client):

helm/clients/image_generation/dalle2_client.py

@@ -15,7 +15,7 @@ from helm.common.tokenization_request import (
 )
 from helm.clients.moderation_api_client import ModerationAPIClient
 from helm.clients.client import Client, CachingClient
-from .image_generation_client_utils import get_single_image_multimedia_object
+from helm.clients.image_generation.image_generation_client_utils import get_single_image_multimedia_object
 
 try:
     import openai

helm/clients/image_generation/dalle3_client.py

@@ -7,8 +7,8 @@ from helm.common.optional_dependencies import handle_module_not_found_error
 from helm.common.request import Request, RequestResult, GeneratedOutput, wrap_request_time
 from helm.clients.moderation_api_client import ModerationAPIClient
 from helm.clients.client import CachingClient
-from .dalle2_client import DALLE2Client
-from .image_generation_client_utils import get_single_image_multimedia_object
+from helm.clients.image_generation.dalle2_client import DALLE2Client
+from helm.clients.image_generation.image_generation_client_utils import get_single_image_multimedia_object
 
 try:
     import openai

helm/clients/image_generation/dalle_mini/__init__.py

@@ -1,3 +1,3 @@
 __version__ = "0.1.4"
 
-from .model import DalleBart, DalleBartProcessor
+from helm.clients.image_generation.dalle_mini.model import DalleBart, DalleBartProcessor

helm/clients/image_generation/dalle_mini/data.py

@@ -6,7 +6,7 @@ from pathlib import Path
 import numpy as np
 from datasets import Dataset, load_dataset
 
-from .model.text import TextNormalizer
+from helm.clients.image_generation.dalle_mini.model.text import TextNormalizer
 from helm.common.optional_dependencies import handle_module_not_found_error
 
 try:

helm/clients/image_generation/dalle_mini/model/__init__.py

@@ -1,5 +1,5 @@
-from .configuration import DalleBartConfig
-from .modeling import DalleBart
-from .partitions import set_partitions
-from .processor import DalleBartProcessor
-from .tokenizer import DalleBartTokenizer
+from helm.clients.image_generation.dalle_mini.model.configuration import DalleBartConfig
+from helm.clients.image_generation.dalle_mini.model.modeling import DalleBart
+from helm.clients.image_generation.dalle_mini.model.partitions import set_partitions
+from helm.clients.image_generation.dalle_mini.model.processor import DalleBartProcessor
+from helm.clients.image_generation.dalle_mini.model.tokenizer import DalleBartTokenizer

helm/clients/image_generation/dalle_mini/model/configuration.py

@@ -18,7 +18,7 @@ import warnings
 from transformers.configuration_utils import PretrainedConfig
 from transformers.utils import logging
 
-from .utils import PretrainedFromWandbMixin
+from helm.clients.image_generation.dalle_mini.model.utils import PretrainedFromWandbMixin
 
 logger = logging.get_logger(__name__)
 

helm/clients/image_generation/dalle_mini/model/modeling.py

@@ -35,8 +35,8 @@ from transformers.utils import ModelOutput, logging
 from transformers.generation.configuration_utils import GenerationConfig
 
 from helm.common.optional_dependencies import handle_module_not_found_error
-from .configuration import DalleBartConfig
-from .utils import PretrainedFromWandbMixin
+from helm.clients.image_generation.dalle_mini.model.configuration import DalleBartConfig
+from helm.clients.image_generation.dalle_mini.model.utils import PretrainedFromWandbMixin
 
 try:
     import flax

helm/clients/image_generation/dalle_mini/model/processor.py

@@ -2,10 +2,10 @@
 
 from typing import List
 
-from .configuration import DalleBartConfig
-from .text import TextNormalizer
-from .tokenizer import DalleBartTokenizer
-from .utils import PretrainedFromWandbMixin
+from helm.clients.image_generation.dalle_mini.model.configuration import DalleBartConfig
+from helm.clients.image_generation.dalle_mini.model.text import TextNormalizer
+from helm.clients.image_generation.dalle_mini.model.tokenizer import DalleBartTokenizer
+from helm.clients.image_generation.dalle_mini.model.utils import PretrainedFromWandbMixin
 from helm.common.optional_dependencies import handle_module_not_found_error
 
 

helm/clients/image_generation/dalle_mini/model/tokenizer.py

@@ -2,7 +2,7 @@
 
 from transformers import BartTokenizerFast
 
-from .utils import PretrainedFromWandbMixin
+from helm.clients.image_generation.dalle_mini.model.utils import PretrainedFromWandbMixin
 
 
 class DalleBartTokenizer(PretrainedFromWandbMixin, BartTokenizerFast):

helm/clients/image_generation/dalle_mini/vqgan_jax/__init__.py

@@ -1 +1 @@
-from . import *
+from helm.clients.image_generation.dalle_mini.vqgan_jax import *

helm/clients/image_generation/dalle_mini/vqgan_jax/convert_pt_model_to_jax.py

@@ -2,8 +2,8 @@ import re
 
 import torch
 
-from .modeling_flax_vqgan import VQModel
-from .configuration_vqgan import VQGANConfig
+from helm.clients.image_generation.dalle_mini.vqgan_jax.modeling_flax_vqgan import VQModel
+from helm.clients.image_generation.dalle_mini.vqgan_jax.configuration_vqgan import VQGANConfig
 from helm.common.optional_dependencies import handle_module_not_found_error
 
 try:

helm/clients/image_generation/dalle_mini/vqgan_jax/modeling_flax_vqgan.py

@@ -6,7 +6,7 @@ import math
 
 from transformers.modeling_flax_utils import FlaxPreTrainedModel
 
-from .configuration_vqgan import VQGANConfig
+from helm.clients.image_generation.dalle_mini.vqgan_jax.configuration_vqgan import VQGANConfig
 from helm.common.optional_dependencies import handle_module_not_found_error
 
 try:

helm/clients/image_generation/dalle_mini_client.py

@@ -15,7 +15,7 @@ from helm.common.tokenization_request import (
     TokenizationRequestResult,
 )
 from helm.clients.client import Client, CachingClient
-from .image_generation_client_utils import get_single_image_multimedia_object
+from helm.clients.image_generation.image_generation_client_utils import get_single_image_multimedia_object
 
 
 class DALLEMiniClient(Client):

helm/clients/image_generation/deep_floyd_client.py

@@ -9,7 +9,7 @@ from helm.common.tokenization_request import (
     DecodeRequestResult,
 )
 from helm.clients.client import Client, CachingClient
-from .image_generation_client_utils import get_single_image_multimedia_object
+from helm.clients.image_generation.image_generation_client_utils import get_single_image_multimedia_object
 
 
 class DeepFloydClient(Client):

helm/clients/image_generation/huggingface_diffusers_client.py

@@ -17,7 +17,7 @@ from helm.common.tokenization_request import (
     TokenizationRequestResult,
 )
 from helm.clients.client import Client, CachingClient
-from .image_generation_client_utils import get_single_image_multimedia_object
+from helm.clients.image_generation.image_generation_client_utils import get_single_image_multimedia_object
 
 
 _models_lock: Lock = Lock()

helm/clients/image_generation/lexica_client.py

@@ -14,7 +14,7 @@ from helm.common.tokenization_request import (
     DecodeRequestResult,
 )
 from helm.clients.client import Client, CachingClient
-from .image_generation_client_utils import get_single_image_multimedia_object
+from helm.clients.image_generation.image_generation_client_utils import get_single_image_multimedia_object
 
 
 class LexicaClient(Client):

helm/clients/image_generation/mindalle/models/__init__.py

@@ -11,12 +11,12 @@ from typing import Optional, Tuple
 from torch.cuda.amp import autocast
 from torch.optim.lr_scheduler import CosineAnnealingLR
 from torch.nn import functional as F
-from .stage1.vqgan import VQGAN
-from .stage2.transformer import Transformer1d, iGPT
-from .. import utils
-from ..utils.config import get_base_config
-from ..utils.sampling import sampling, sampling_igpt
-from .tokenizer import build_tokenizer
+from helm.clients.image_generation.mindalle.models.stage1.vqgan import VQGAN
+from helm.clients.image_generation.mindalle.models.stage2.transformer import Transformer1d, iGPT
+from helm.clients.image_generation.mindalle import utils
+from helm.clients.image_generation.mindalle.utils.config import get_base_config
+from helm.clients.image_generation.mindalle.utils.sampling import sampling, sampling_igpt
+from helm.clients.image_generation.mindalle.models.tokenizer import build_tokenizer
 from helm.common.optional_dependencies import handle_module_not_found_error
 
 

helm/clients/image_generation/mindalle/models/stage1/vqgan.py

@@ -7,7 +7,7 @@ import torch
 import torch.nn as nn
 from typing import List, Tuple, Optional
 
-from .layers import Encoder, Decoder
+from helm.clients.image_generation.mindalle.models.stage1.layers import Encoder, Decoder
 from helm.common.optional_dependencies import handle_module_not_found_error
 
 

helm/clients/image_generation/mindalle/models/stage2/transformer.py

@@ -11,7 +11,7 @@ import torch
 import torch.nn as nn
 from typing import Optional, Tuple, List
 from torch.cuda.amp import autocast
-from .layers import Block
+from helm.clients.image_generation.mindalle.models.stage2.layers import Block
 
 from helm.common.optional_dependencies import handle_module_not_found_error
 

helm/clients/image_generation/mindalle/utils/__init__.py

@@ -1,3 +1,3 @@
-from .utils import *
-from .config import *
-from .sampling import *
+from helm.clients.image_generation.mindalle.utils.utils import *
+from helm.clients.image_generation.mindalle.utils.config import *
+from helm.clients.image_generation.mindalle.utils.sampling import *

helm/clients/image_generation/mindalle_client.py

@@ -15,7 +15,7 @@ from helm.common.tokenization_request import (
     TokenizationRequestResult,
 )
 from helm.clients.client import Client, CachingClient
-from .image_generation_client_utils import get_single_image_multimedia_object
+from helm.clients.image_generation.image_generation_client_utils import get_single_image_multimedia_object
 
 try:
     from PIL import Image

helm/clients/image_generation/together_image_generation_client.py

@@ -13,7 +13,7 @@ from helm.common.tokenization_request import (
 )
 
 from helm.clients.client import CachingClient, Client
-from .image_generation_client_utils import get_single_image_multimedia_object
+from helm.clients.image_generation.image_generation_client_utils import get_single_image_multimedia_object
 
 
 class TogetherImageGenerationClient(Client):

helm/clients/lit_gpt_client.py

@@ -12,8 +12,8 @@ from helm.common.optional_dependencies import OptionalDependencyNotInstalled
 from helm.common.request import Request, RequestResult, GeneratedOutput, Token
 from helm.tokenizers.tokenizer import Tokenizer
 
-from .client import CachingClient
-from .lit_gpt_generate import generate  # type: ignore
+from helm.clients.client import CachingClient
+from helm.clients.lit_gpt_generate import generate  # type: ignore
 
 try:
     import lightning as L