codekingpro/portable-devtools

1import importlib.util2import logging3from typing import Any, Callable, List, Mapping, Optional4 5from langchain_core.callbacks import CallbackManagerForLLMRun6from pydantic import ConfigDict7 8from langchain_community.llms.self_hosted import SelfHostedPipeline9from langchain_community.llms.utils import enforce_stop_tokens10 11DEFAULT_MODEL_ID = "gpt2"12DEFAULT_TASK = "text-generation"13VALID_TASKS = ("text2text-generation", "text-generation", "summarization")14 15logger = logging.getLogger(__name__)16 17 18def _generate_text(19    pipeline: Any,20    prompt: str,21    *args: Any,22    stop: Optional[List[str]] = None,23    **kwargs: Any,24) -> str:25    """Inference function to send to the remote hardware.26 27    Accepts a Hugging Face pipeline (or more likely,28    a key pointing to such a pipeline on the cluster's object store)29    and returns generated text.30    """31    response = pipeline(prompt, *args, **kwargs)32    if pipeline.task == "text-generation":33        # Text generation return includes the starter text.34        text = response[0]["generated_text"][len(prompt) :]35    elif pipeline.task == "text2text-generation":36        text = response[0]["generated_text"]37    elif pipeline.task == "summarization":38        text = response[0]["summary_text"]39    else:40        raise ValueError(41            f"Got invalid task {pipeline.task}, "42            f"currently only {VALID_TASKS} are supported"43        )44    if stop is not None:45        text = enforce_stop_tokens(text, stop)46    return text47 48 49def _load_transformer(50    model_id: str = DEFAULT_MODEL_ID,51    task: str = DEFAULT_TASK,52    device: int = 0,53    model_kwargs: Optional[dict] = None,54) -> Any:55    """Inference function to send to the remote hardware.56 57    Accepts a huggingface model_id and returns a pipeline for the task.58    """59    from transformers import AutoModelForCausalLM, AutoModelForSeq2SeqLM, AutoTokenizer60    from transformers import pipeline as hf_pipeline61 62    _model_kwargs = model_kwargs or {}63    tokenizer = AutoTokenizer.from_pretrained(model_id, **_model_kwargs)64 65    try:66        if task == "text-generation":67            model = AutoModelForCausalLM.from_pretrained(model_id, **_model_kwargs)68        elif task in ("text2text-generation", "summarization"):69            model = AutoModelForSeq2SeqLM.from_pretrained(model_id, **_model_kwargs)70        else:71            raise ValueError(72                f"Got invalid task {task}, currently only {VALID_TASKS} are supported"73            )74    except ImportError as e:75        raise ImportError(76            f"Could not load the {task} model due to missing dependencies."77        ) from e78 79    if importlib.util.find_spec("torch") is not None:80        import torch81 82        cuda_device_count = torch.cuda.device_count()83        if device < -1 or (device >= cuda_device_count):84            raise ValueError(85                f"Got device=={device}, "86                f"device is required to be within [-1, {cuda_device_count})"87            )88        if device < 0 and cuda_device_count > 0:89            logger.warning(90                "Device has %d GPUs available. "91                "Provide device={deviceId} to `from_model_id` to use available"92                "GPUs for execution. deviceId is -1 for CPU and "93                "can be a positive integer associated with CUDA device id.",94                cuda_device_count,95            )96 97    pipeline = hf_pipeline(98        task=task,99        model=model,100        tokenizer=tokenizer,101        device=device,102        model_kwargs=_model_kwargs,103    )104    if pipeline.task not in VALID_TASKS:105        raise ValueError(106            f"Got invalid task {pipeline.task}, "107            f"currently only {VALID_TASKS} are supported"108        )109    return pipeline110 111 112class SelfHostedHuggingFaceLLM(SelfHostedPipeline):113    """HuggingFace Pipeline API to run on self-hosted remote hardware.114 115    Supported hardware includes auto-launched instances on AWS, GCP, Azure,116    and Lambda, as well as servers specified117    by IP address and SSH credentials (such as on-prem, or another cloud118    like Paperspace, Coreweave, etc.).119 120    To use, you should have the ``runhouse`` python package installed.121 122    Only supports `text-generation`, `text2text-generation` and `summarization` for now.123 124    Example using from_model_id:125        .. code-block:: python126 127            from langchain_community.llms import SelfHostedHuggingFaceLLM128            import runhouse as rh129            gpu = rh.cluster(name="rh-a10x", instance_type="A100:1")130            hf = SelfHostedHuggingFaceLLM(131                model_id="google/flan-t5-large", task="text2text-generation",132                hardware=gpu133            )134    Example passing fn that generates a pipeline (bc the pipeline is not serializable):135        .. code-block:: python136 137            from langchain_community.llms import SelfHostedHuggingFaceLLM138            from transformers import AutoModelForCausalLM, AutoTokenizer, pipeline139            import runhouse as rh140 141            def get_pipeline():142                model_id = "gpt2"143                tokenizer = AutoTokenizer.from_pretrained(model_id)144                model = AutoModelForCausalLM.from_pretrained(model_id)145                pipe = pipeline(146                    "text-generation", model=model, tokenizer=tokenizer147                )148                return pipe149            hf = SelfHostedHuggingFaceLLM(150                model_load_fn=get_pipeline, model_id="gpt2", hardware=gpu)151    """152 153    model_id: str = DEFAULT_MODEL_ID154    """Hugging Face model_id to load the model."""155    task: str = DEFAULT_TASK156    """Hugging Face task ("text-generation", "text2text-generation" or157    "summarization")."""158    device: int = 0159    """Device to use for inference. -1 for CPU, 0 for GPU, 1 for second GPU, etc."""160    model_kwargs: Optional[dict] = None161    """Keyword arguments to pass to the model."""162    hardware: Any = None163    """Remote hardware to send the inference function to."""164    model_reqs: List[str] = ["./", "transformers", "torch"]165    """Requirements to install on hardware to inference the model."""166    model_load_fn: Callable = _load_transformer167    """Function to load the model remotely on the server."""168    inference_fn: Callable = _generate_text  #: :meta private:169    """Inference function to send to the remote hardware."""170 171    model_config = ConfigDict(172        extra="forbid",173    )174 175    def __init__(self, **kwargs: Any):176        """Construct the pipeline remotely using an auxiliary function.177 178        The load function needs to be importable to be imported179        and run on the server, i.e. in a module and not a REPL or closure.180        Then, initialize the remote inference function.181        """182        load_fn_kwargs = {183            "model_id": kwargs.get("model_id", DEFAULT_MODEL_ID),184            "task": kwargs.get("task", DEFAULT_TASK),185            "device": kwargs.get("device", 0),186            "model_kwargs": kwargs.get("model_kwargs", None),187        }188        super().__init__(load_fn_kwargs=load_fn_kwargs, **kwargs)189 190    @property191    def _identifying_params(self) -> Mapping[str, Any]:192        """Get the identifying parameters."""193        return {194            **{"model_id": self.model_id},195            **{"model_kwargs": self.model_kwargs},196        }197 198    @property199    def _llm_type(self) -> str:200        return "selfhosted_huggingface_pipeline"201 202    def _call(203        self,204        prompt: str,205        stop: Optional[List[str]] = None,206        run_manager: Optional[CallbackManagerForLLMRun] = None,207        **kwargs: Any,208    ) -> str:209        return self.client(210            pipeline=self.pipeline_ref, prompt=prompt, stop=stop, **kwargs211        )212