ai-parrot 0.8.3__cp311-cp311-manylinux2014_x86_64.manylinux_2_17_x86_64.whl

parrot/llms/hf.py

@@ -0,0 +1,45 @@
+from langchain_community.llms import HuggingFacePipeline  # pylint: disable=import-error, E0611
+from transformers import AutoModelForCausalLM, AutoTokenizer
+from .abstract import AbstractLLM
+
+class HuggingFace(AbstractLLM):
+    """HuggingFace.
+
+    Load a LLM (Language Model) from HuggingFace Hub.
+
+    Only supports text-generation, text2text-generation, summarization and translation for now.
+
+    Returns:
+        _type_: an instance of HuggingFace LLM Model.
+    """
+    model: str = "databricks/dolly-v2-3b"
+    embed_model: str = None
+    max_tokens: int = 1024
+    supported_models: list = [
+        "databricks/dolly-v2-3b",
+        "gpt2",
+        "bigscience/bloom-1b7",
+        "meta-llama/Llama-2-7b-hf"
+    ]
+
+    def __init__(self, *args, **kwargs):
+        self.batch_size = kwargs.get('batch_size', 4)
+        super().__init__(*args, **kwargs)
+        self._tokenizer = AutoTokenizer.from_pretrained(self.model, chunk_size=self.max_tokens)
+        self._model = AutoModelForCausalLM.from_pretrained(self.model, trust_remote_code=True)
+        self._llm = HuggingFacePipeline.from_model_id(
+            model_id=self.model,
+            task=self.task,
+            device_map='auto',
+            batch_size=self.batch_size,
+            model_kwargs={
+                "max_length": self.max_tokens,
+                "trust_remote_code": True
+            },
+            pipeline_kwargs={
+                "temperature": self.temperature,
+                "repetition_penalty":1.1,
+                "max_new_tokens": self.max_tokens,
+                **self.args
+            }
+        )

parrot/llms/openai.py

@@ -0,0 +1,61 @@
+from langchain_openai import (  # pylint: disable=E0401, E0611
+    OpenAI,
+    ChatOpenAI,
+)
+from navconfig import config
+from navconfig.logging import logging
+from .abstract import AbstractLLM
+
+
+logging.getLogger(name='openai').setLevel(logging.WARNING)
+logging.getLogger(name='httpcore').setLevel(logging.WARNING)
+logging.getLogger(name='httpx').setLevel(logging.WARNING)
+
+class OpenAILLM(AbstractLLM):
+    """OpenAI.
+    Interact with OpenAI Language Model.
+
+    Returns:
+        _type_: an instance of OpenAI LLM Model.
+    """
+    model: str = "gpt-4-turbo"
+    max_tokens: int = 8192
+    top_k: float = 40
+    top_p: float = 1.0
+    supported_models: list = [
+        "gpt-4.1",
+        "gpt-4o-mini",
+        'gpt-4.1-2025-04-14',
+        'o4-mini-2025-04-16',
+        "o3-2025-04-16",
+        'gpt-4-turbo',
+        'gpt-4o',
+        'gpt-3.5-turbo',
+        'gpt-3.5-turbo-instruct',
+        'dall-e-3'
+        'tts-1',
+    ]
+
+    def __init__(self, *args, use_chat: bool = False, **kwargs):
+        self.model_type = kwargs.get("model_type", "text")
+        super().__init__(*args, **kwargs)
+        self.model = kwargs.get("model", "davinci")
+        self._api_key = kwargs.pop('api_key', config.get('OPENAI_API_KEY'))
+        organization = config.get("OPENAI_ORGANIZATION")
+        if use_chat:
+            base_llm = ChatOpenAI
+        else:
+            base_llm = OpenAI
+        args = {
+            "api_key": self._api_key,
+            "organization": organization,
+            "temperature": self.temperature,
+            "max_tokens": self.max_tokens,
+            "max_retries": 4,
+            "top_p": self.top_p,
+            "verbose": True
+        }
+        self._llm = base_llm(
+            model_name=self.model,
+            **args
+        )

parrot/llms/pipes.py

@@ -0,0 +1,114 @@
+import torch
+from langchain_community.llms import HuggingFacePipeline  # pylint: disable=import-error, E0611
+from transformers import (
+    AutoModelForCausalLM,
+    AutoProcessor,
+    LlavaForConditionalGeneration,
+    AutoTokenizer,
+    GenerationConfig,
+    pipeline
+)
+from .abstract import AbstractLLM
+
+
+class PipelineLLM(AbstractLLM):
+    """PipelineLLM.
+
+    Load a LLM (Language Model) from HuggingFace Hub.
+
+    Returns:
+        _type_: an instance of HuggingFace LLM Model.
+    """
+    model: str = "databricks/dolly-v2-3b"
+    embed_model: str = None
+    max_tokens: int = 1024
+    supported_models: list = [
+        "databricks/dolly-v2-3b",
+        "gpt2",
+        "bigscience/bloom-1b7",
+        "meta-llama/Llama-2-7b-hf",
+        'llava-hf/llava-1.5-7b-hf'
+    ]
+
+    def __init__(self, *args, **kwargs):
+        self.batch_size = kwargs.get('batch_size', 4)
+        self.use_llava: bool = kwargs.get('use_llava', False)
+        self.model_args = kwargs.get('model_args', {})
+        super().__init__(*args, **kwargs)
+        dtype = kwargs.get('dtype', 'float16')
+        if dtype == 'bfloat16':
+            torch_dtype = torch.bfloat16
+        if dtype == 'float16':
+            torch_dtype = torch.float16
+        elif dtype == 'float32':
+            torch_dtype = torch.float32
+        elif dtype == 'float8':
+            torch_dtype = torch.float8
+        else:
+            torch_dtype = "auto"
+        use_fast = kwargs.get('use_fast', True)
+        if self.use_llava is False:
+            self.tokenizer = AutoTokenizer.from_pretrained(
+                self.model,
+                chunk_size=self.max_tokens
+            )
+            self._model = AutoModelForCausalLM.from_pretrained(
+                self.model,
+                device_map="auto",
+                torch_dtype=torch_dtype,
+                trust_remote_code=True,
+            )
+            config = GenerationConfig(
+                do_sample=True,
+                temperature=self.temperature,
+                max_new_tokens=self.max_tokens,
+                top_p=self.top_p,
+                top_k=self.top_k,
+                repetition_penalty=1.15,
+            )
+            self._pipe = pipeline(
+                task=self.task,
+                model=self._model,
+                tokenizer=self.tokenizer,
+                return_full_text=True,
+                use_fast=use_fast,
+                device_map='auto',
+                batch_size=self.batch_size,
+                generation_config=config,
+                pad_token_id = 50256,
+                framework="pt"
+            )
+        else:
+            self._model = LlavaForConditionalGeneration.from_pretrained(
+                self.model,
+                device_map="auto",
+                torch_dtype=torch_dtype,
+                trust_remote_code=True,
+                low_cpu_mem_usage=True,
+            )
+            self.tokenizer = AutoTokenizer.from_pretrained(self.model)
+            processor = AutoProcessor.from_pretrained(self.model)
+            self._pipe = pipeline(
+                task=self.task,
+                model=self._model,
+                tokenizer=self.tokenizer,
+                use_fast=use_fast,
+                device_map='auto',
+                batch_size=self.batch_size,
+                image_processor=processor.image_processor,
+                framework="pt",
+                **self.model_args
+            )
+        self._pipe.tokenizer.pad_token_id = self._pipe.model.config.eos_token_id
+        self._llm = HuggingFacePipeline(
+            model_id=self.model,
+            pipeline=self._pipe,
+            verbose=True
+        )
+
+    def pipe(self, *args, **kwargs):
+        return self._pipe(
+            *args,
+            **kwargs,
+            generate_kwargs={"max_new_tokens": self.max_tokens}
+        )

parrot/llms/vertex.py

@@ -0,0 +1,89 @@
+import os
+from navconfig import config, BASE_DIR
+from google.cloud import aiplatform
+from google.oauth2 import service_account
+from vertexai.preview.vision_models import ImageGenerationModel
+from langchain_google_vertexai import (
+    ChatVertexAI,
+    VertexAI,
+    HarmBlockThreshold,
+    HarmCategory
+)
+from navconfig.logging import logging
+from .abstract import AbstractLLM
+
+logging.getLogger(name='httpcore').setLevel(logging.WARNING)
+logging.getLogger(name='httpx').setLevel(logging.WARNING)
+
+safety_settings = {
+    HarmCategory.HARM_CATEGORY_UNSPECIFIED: HarmBlockThreshold.BLOCK_NONE,
+    HarmCategory.HARM_CATEGORY_SEXUALLY_EXPLICIT: HarmBlockThreshold.BLOCK_NONE,
+    HarmCategory.HARM_CATEGORY_HARASSMENT: HarmBlockThreshold.BLOCK_LOW_AND_ABOVE,
+    HarmCategory.HARM_CATEGORY_DANGEROUS_CONTENT: HarmBlockThreshold.BLOCK_MEDIUM_AND_ABOVE,
+    HarmCategory.HARM_CATEGORY_HATE_SPEECH: HarmBlockThreshold.BLOCK_ONLY_HIGH,
+}
+
+
+class VertexLLM(AbstractLLM):
+    """VertexLLM.
+
+    Interact with VertexAI Language Model.
+
+    Returns:
+        _type_: VertexAI LLM.
+    """
+    model: str = "gemini-1.5-pro"
+    max_tokens: int = 8192
+    top_k: float = 40
+    top_p: float = 0.95
+    supported_models: list = [
+        "gemini-2.5-pro-exp-03-25",
+        "gemini-2.5-pro-preview-03-25",
+        "gemini-2.5-flash-preview-04-17",
+        "gemini-2.0-flash",
+        "gemini-2.0-flash-lite",
+        "gemini-2.0-flash-001",
+        "gemini-1.5-pro",
+        "gemini-1.5-flash-8b",
+        "gemini-1.5-pro-exp-0801",
+        "gemini-1.5-flash-preview-0514",
+        "gemini-1.5-flash-001",
+        "chat-bison@001",
+        "chat-bison@002",
+        "imagen-3.0-generate-002",
+        "gemini-2.0-flash-live-001",
+        "veo-2.0-generate-001"
+    ]
+
+    def __init__(self, *args, use_chat: bool = False, **kwargs):
+        super().__init__(*args, **kwargs)
+        project_id = config.get("VERTEX_PROJECT_ID")
+        region = config.get("VERTEX_REGION")
+        config_file = config.get('GOOGLE_CREDENTIALS_FILE', 'env/google/vertexai.json')
+        config_dir = BASE_DIR.joinpath(config_file)
+        vertex_credentials = service_account.Credentials.from_service_account_file(
+            str(config_dir)
+        )
+        # os.environ["GOOGLE_APPLICATION_CREDENTIALS"] = str(config_dir)
+        args = {
+            "project": project_id,
+            "location": region,
+            "temperature": self.temperature,
+            "max_tokens": self.max_tokens,
+            "max_retries": 4,
+            "top_p": self.top_p,
+            # "top_k": self.top_k,
+            "verbose": True,
+            "credentials": vertex_credentials,
+            # "safety_settings": safety_settings
+        }
+        if use_chat is True:
+            base_llm = ChatVertexAI
+        else:
+            base_llm = VertexAI
+        self._llm = base_llm(
+            model_name=self.model,
+            **args
+        )
+        # LLM
+        self._version_ = aiplatform.__version__

parrot/loaders/__init__.py

@@ -0,0 +1,9 @@
+####
+# Copyright 2023 Jesus Lara.
+# Licensed under the Apache License, Version 2.0 (the "License");
+#
+# Loaders.
+# Open, extract and load data from different sources.
+#####
+from langchain.docstore.document import Document
+from .abstract import AbstractLoader