Fix pypi & black

* Fix pypi & black

Co-authored-by: vatorre <[email protected]>

promptmeteo/__init__.py

@@ -4,4 +4,4 @@
 from .document_classifier import DocumentClassifier
 from .api_generator import APIGenerator
 from .api_formatter import APIFormatter
-from .summarizer import Summarizer
+from .summarizer import Summarizer

promptmeteo/api_formatter.py

@@ -43,7 +43,7 @@
 
 
 class APIFormatter(BaseUnsupervised):
-    """ API Formatter Task.
+    """API Formatter Task.
 
     This class initializes the API Formatter Task to correct and format APIs.
 
@@ -397,6 +397,7 @@ def _add_external_information(api: str, replacements: dict) -> str:
         str
             Updated API YAML string.
         """
+
         def replace_values(orig_dict, replace_dict):
             for k, v in replace_dict.items():
                 if k in orig_dict:

promptmeteo/base.py

@@ -44,7 +44,6 @@
 
 
 class Base(ABC):
-
     """
     Promptmeteo is a tool powered by LLMs, capable of solving NLP tasks such as
     text classification and Named Entity Recognition. Its interface resembles
@@ -347,7 +346,6 @@ def _load_builder(self, **kwargs) -> TaskBuilder:
 
 
 class BaseSupervised(Base):
-
     """
     Base class for supervised training tasks.
     """

promptmeteo/document_classifier.py

@@ -33,7 +33,6 @@
 
 
 class DocumentClassifier(BaseSupervised):
-
     """
     DocumentClassifier Task
 

promptmeteo/document_qa.py

@@ -26,7 +26,6 @@
 
 
 class DocumentQA(BaseUnsupervised):
-
     """
     Question Answering over Documents Task
 

promptmeteo/models/__init__.py

@@ -33,7 +33,6 @@
 
 
 class ModelProvider(str, Enum):
-
     """
     LLM providers currently supported by Promptmeteo
     """
@@ -47,7 +46,6 @@ class ModelProvider(str, Enum):
 
 
 class ModelFactory:
-
     """
     The ModelFactory class is used to create a BaseModel object from the given
     configuration.
@@ -59,7 +57,7 @@ class ModelFactory:
         ModelProvider.PROVIDER_2: HFHubApiLLM,
         ModelProvider.PROVIDER_3: HFPipelineLLM,
         ModelProvider.PROVIDER_3: GoogleVertexAILLM,
-        ModelProvider.PROVIDER_5: BedrockLLM
+        ModelProvider.PROVIDER_5: BedrockLLM,
     }
 
     @classmethod
@@ -90,7 +88,7 @@ def factory_method(
 
         elif model_provider_name == ModelProvider.PROVIDER_4.value:
             model_cls = GoogleVertexAILLM
-            
+
         elif model_provider_name == ModelProvider.PROVIDER_5.value:
             model_cls = BedrockLLM
 

promptmeteo/models/azure_openai.py

@@ -33,7 +33,6 @@
 
 
 class ModelTypes(str, Enum):
-
     """
     Enum of available model types.
     """
@@ -54,13 +53,11 @@ def has_value(
 
 
 class ModelEnum(Enum):
-
     """
     Model types with their parameters.
     """
 
     class GPT35TurboInstruct:
-
         """
         Default parameters for TextDavinci003 model.
         """
@@ -76,7 +73,6 @@ class GPT35TurboInstruct:
         }
 
     class GPT35Turbo:
-
         """
         Default parameters for GPT35Turbo model.
         """
@@ -93,7 +89,6 @@ class GPT35Turbo:
 
 
 class AzureOpenAILLM(BaseModel):
-
     """
     OpenAI LLM model.
     """

promptmeteo/models/base.py

@@ -21,20 +21,21 @@
 #  THE SOFTWARE.
 
 from abc import ABC
+from typing import Optional
+
 from langchain.llms.base import BaseLLM
 from langchain.schema import HumanMessage
 from langchain.embeddings.base import Embeddings
 
 
 class BaseModel(ABC):
-
     """
     Model Interface.
     """
 
-    def __init__(self):
-        self._llm: BaseLLM = None
-        self._embeddings: Embeddings = None
+    def __init__(self, **kwargs):
+        self._llm: Optional[BaseLLM] = kwargs.get("llm", None)
+        self._embeddings: Optional[Embeddings] = kwargs.get("embeddings", None)
 
     @property
     def llm(
@@ -59,10 +60,10 @@ def run(
         """
 
         try:
-            return self._llm(prompt=sample)
+            return self.llm(prompt=sample)
 
         except TypeError:
-            return self._llm([HumanMessage(content=sample)]).content
+            return self.llm([HumanMessage(content=sample)]).content
 
         except Exception as error:
             raise RuntimeError(

promptmeteo/models/bedrock.py

@@ -32,7 +32,6 @@
 
 
 class ModelTypes(str, Enum):
-
     """
     Enum of available model types.
     """
@@ -52,13 +51,11 @@ def has_value(
 
 
 class ModelEnum(Enum):
-
     """
     Model Parameters.
     """
 
     class AnthropicClaudeV2:
-
         """
         Default parameters for Anthropic Claude V2
         """
@@ -67,16 +64,15 @@ class AnthropicClaudeV2:
         embedding = HuggingFaceEmbeddings
         model_task: str = "text2text-generation"
         params: dict = {
-            'max_tokens_to_sample': 2048,
-            'temperature': 0.3,
-            'top_k': 250,
-            'top_p': 0.999,
-            'stop_sequences': ['Human:']
+            "max_tokens_to_sample": 2048,
+            "temperature": 0.3,
+            "top_k": 250,
+            "top_p": 0.999,
+            "stop_sequences": ["Human:"],
         }
 
 
 class BedrockLLM(BaseModel):
-
     """
     Bedrock LLM model.
     """
@@ -86,7 +82,7 @@ def __init__(
         model_name: Optional[str] = "",
         model_params: Optional[Dict] = None,
         model_provider_token: Optional[str] = "",
-        **kwargs
+        **kwargs,
     ) -> None:
         """
         Make predictions using a model from OpenAI.
@@ -98,7 +94,7 @@ def __init__(
                 f"`model_name`={model_name} not in supported model names: "
                 f"{[i.name for i in ModelTypes]}"
             )
-        self.boto3_bedrock = boto3.client('bedrock-runtime', **kwargs)
+        self.boto3_bedrock = boto3.client("bedrock-runtime", **kwargs)
         super(BedrockLLM, self).__init__()
 
         # Model name
@@ -117,7 +113,7 @@ def __init__(
         self._llm = ModelEnum[model].value.client(
             model_id=model_name,
             model_kwargs=self.model_params,
-            client =  self.boto3_bedrock
+            client=self.boto3_bedrock,
         )
 
         embedding_name = "sentence-transformers/all-MiniLM-L6-v2"

promptmeteo/models/fake_llm.py

@@ -35,7 +35,6 @@
 
 
 class FakeStaticLLM(LLM):
-
     """
     Fake Static LLM wrapper for testing purposes.
     """
@@ -82,7 +81,6 @@ async def _acall(
 
 
 class FakePromptCopyLLM(LLM):
-
     """
     Fake Prompt Copy LLM wrapper for testing purposes.
     """
@@ -127,7 +125,6 @@ async def _acall(
 
 
 class FakeListLLM(LLM):
-
     """
     Fake LLM wrapper for testing purposes.
     """
@@ -178,7 +175,6 @@ async def _acall(
 
 
 class ModelTypes(Enum):
-
     """
     FakeLLM Model Types.
     """

promptmeteo/models/google_vertexai.py

@@ -32,7 +32,6 @@
 
 
 class ModelTypes(str, Enum):
-
     """
     Enum of available model types.
     """
@@ -55,13 +54,11 @@ def has_value(
 
 
 class ModelEnum(Enum):
-
     """
     Model types with their parameters.
     """
 
     class TextBison001:
-
         """
         Default parameters for text-bison model.
         """
@@ -70,7 +67,6 @@ class TextBison001:
         model_kwargs = {"temperature": 0.4, "max_tokens": 256, "max_retries": 3}
 
     class TextBison:
-
         """
         Default parameters for text-bison model in their latest version
         """
@@ -79,7 +75,6 @@ class TextBison:
         model_kwargs = {"temperature": 0.4, "max_tokens": 256, "max_retries": 3}
 
     class TextBison32k:
-
         """
         Default parameters for text-bison-32 model in their latest version
         """
@@ -89,7 +84,6 @@ class TextBison32k:
 
 
 class GoogleVertexAILLM(BaseModel):
-
     """
     Google VertexAI LLM model.
     """
@@ -129,4 +123,4 @@ def __init__(
         # Model Parameters
         if not model_params:
             model_params = ModelEnum[model].value
-        self.model_params = model_params
+        self.model_params = model_params

promptmeteo/models/hf_hub_api.py

@@ -30,7 +30,6 @@
 
 
 class ModelTypes(str, Enum):
-
     """
     Enum of available model types.
     """
@@ -50,13 +49,11 @@ def has_value(
 
 
 class ModelEnum(Enum):
-
     """
     Model Parameters Enum
     """
 
     class FlanT5Xxl:
-
         """
         Flan-t5-xxl default params
         """
@@ -66,7 +63,6 @@ class FlanT5Xxl:
 
 
 class HFHubApiLLM(BaseModel):
-
     """
     HuggingFace API call.
     """

promptmeteo/models/hf_pipeline.py

@@ -31,7 +31,6 @@
 
 
 class ModelTypes(str, Enum):
-
     """
     Enum of available model types.
     """
@@ -48,13 +47,11 @@ def has_value(cls, value):
 
 
 class ModelParams(Enum):
-
     """
     Model Parameters.
     """
 
     class MODEL_1:
-
         """
         Parameters Model 1.
         """
@@ -65,7 +62,6 @@ class MODEL_1:
 
 
 class HFPipelineLLM(BaseModel):
-
     """
     HuggingFace Local Pipeline.
     """