codekingpro/portable-devtools

1import base642import json3import logging4import subprocess5import textwrap6import time7from typing import Any, Dict, List, Mapping, Optional8 9import requests10from langchain_core.callbacks import CallbackManagerForLLMRun11from langchain_core.language_models.llms import LLM12from langchain_core.utils import get_from_dict_or_env, pre_init13from langchain_core.utils.pydantic import get_fields14from pydantic import ConfigDict, Field, model_validator15 16logger = logging.getLogger(__name__)17 18DEFAULT_NUM_TRIES = 1019DEFAULT_SLEEP_TIME = 420 21 22class Beam(LLM):23    """Beam API for gpt2 large language model.24 25    To use, you should have the ``beam-sdk`` python package installed,26    and the environment variable ``BEAM_CLIENT_ID`` set with your client id27    and ``BEAM_CLIENT_SECRET`` set with your client secret. Information on how28    to get this is available here: https://docs.beam.cloud/account/api-keys.29 30    The wrapper can then be called as follows, where the name, cpu, memory, gpu,31    python version, and python packages can be updated accordingly. Once deployed,32    the instance can be called.33 34    Example:35        .. code-block:: python36 37            llm = Beam(model_name="gpt2",38                name="langchain-gpt2",39                cpu=8,40                memory="32Gi",41                gpu="A10G",42                python_version="python3.8",43                python_packages=[44                    "diffusers[torch]>=0.10",45                    "transformers",46                    "torch",47                    "pillow",48                    "accelerate",49                    "safetensors",50                    "xformers",],51                max_length=50)52            llm._deploy()53            call_result = llm._call(input)54 55    """56 57    model_name: str = ""58    name: str = ""59    cpu: str = ""60    memory: str = ""61    gpu: str = ""62    python_version: str = ""63    python_packages: List[str] = []64    max_length: str = ""65    url: str = ""66    """model endpoint to use"""67 68    model_kwargs: Dict[str, Any] = Field(default_factory=dict)69    """Holds any model parameters valid for `create` call not70    explicitly specified."""71 72    beam_client_id: str = ""73    beam_client_secret: str = ""74    app_id: Optional[str] = None75 76    model_config = ConfigDict(77        extra="forbid",78    )79 80    @model_validator(mode="before")81    @classmethod82    def build_extra(cls, values: Dict[str, Any]) -> Any:83        """Build extra kwargs from additional params that were passed in."""84        all_required_field_names = {field.alias for field in get_fields(cls).values()}85 86        extra = values.get("model_kwargs", {})87        for field_name in list(values):88            if field_name not in all_required_field_names:89                if field_name in extra:90                    raise ValueError(f"Found {field_name} supplied twice.")91                logger.warning(92                    f"""{field_name} was transferred to model_kwargs.93                    Please confirm that {field_name} is what you intended."""94                )95                extra[field_name] = values.pop(field_name)96        values["model_kwargs"] = extra97        return values98 99    @pre_init100    def validate_environment(cls, values: Dict) -> Dict:101        """Validate that api key and python package exists in environment."""102        beam_client_id = get_from_dict_or_env(103            values, "beam_client_id", "BEAM_CLIENT_ID"104        )105        beam_client_secret = get_from_dict_or_env(106            values, "beam_client_secret", "BEAM_CLIENT_SECRET"107        )108        values["beam_client_id"] = beam_client_id109        values["beam_client_secret"] = beam_client_secret110        return values111 112    @property113    def _identifying_params(self) -> Mapping[str, Any]:114        """Get the identifying parameters."""115        return {116            "model_name": self.model_name,117            "name": self.name,118            "cpu": self.cpu,119            "memory": self.memory,120            "gpu": self.gpu,121            "python_version": self.python_version,122            "python_packages": self.python_packages,123            "max_length": self.max_length,124            "model_kwargs": self.model_kwargs,125        }126 127    @property128    def _llm_type(self) -> str:129        """Return type of llm."""130        return "beam"131 132    def app_creation(self) -> None:133        """Creates a Python file which will contain your Beam app definition."""134        script = textwrap.dedent(135            """\136        import beam137 138        # The environment your code will run on139        app = beam.App(140            name="{name}",141            cpu={cpu},142            memory="{memory}",143            gpu="{gpu}",144            python_version="{python_version}",145            python_packages={python_packages},146        )147 148        app.Trigger.RestAPI(149            inputs={{"prompt": beam.Types.String(), "max_length": beam.Types.String()}},150            outputs={{"text": beam.Types.String()}},151            handler="run.py:beam_langchain",152        )153 154        """155        )156 157        script_name = "app.py"158        with open(script_name, "w") as file:159            file.write(160                script.format(161                    name=self.name,162                    cpu=self.cpu,163                    memory=self.memory,164                    gpu=self.gpu,165                    python_version=self.python_version,166                    python_packages=self.python_packages,167                )168            )169 170    def run_creation(self) -> None:171        """Creates a Python file which will be deployed on beam."""172        script = textwrap.dedent(173            """174        import os175        import transformers176        from transformers import GPT2LMHeadModel, GPT2Tokenizer177 178        model_name = "{model_name}"179 180        def beam_langchain(**inputs):181            prompt = inputs["prompt"]182            length = inputs["max_length"]183 184            tokenizer = GPT2Tokenizer.from_pretrained(model_name)185            model = GPT2LMHeadModel.from_pretrained(model_name)186            encodedPrompt = tokenizer.encode(prompt, return_tensors='pt')187            outputs = model.generate(encodedPrompt, max_length=int(length),188              do_sample=True, pad_token_id=tokenizer.eos_token_id)189            output = tokenizer.decode(outputs[0], skip_special_tokens=True)190 191            print(output)  # noqa: T201192            return {{"text": output}}193 194        """195        )196 197        script_name = "run.py"198        with open(script_name, "w") as file:199            file.write(script.format(model_name=self.model_name))200 201    def _deploy(self) -> str:202        """Call to Beam."""203        try:204            import beam205 206            if beam.__path__ == "":207                raise ImportError208        except ImportError:209            raise ImportError(210                "Could not import beam python package. "211                "Please install it with `curl "212                "https://raw.githubusercontent.com/slai-labs"213                "/get-beam/main/get-beam.sh -sSfL | sh`."214            )215        self.app_creation()216        self.run_creation()217 218        process = subprocess.run(219            "beam deploy app.py", shell=True, capture_output=True, text=True220        )221 222        if process.returncode == 0:223            output = process.stdout224            logger.info(output)225            lines = output.split("\n")226 227            for line in lines:228                if line.startswith(" i  Send requests to: https://apps.beam.cloud/"):229                    self.app_id = line.split("/")[-1]230                    self.url = line.split(":")[1].strip()231                    return self.app_id232 233            raise ValueError(234                f"""Failed to retrieve the appID from the deployment output.235                Deployment output: {output}"""236            )237        else:238            raise ValueError(f"Deployment failed. Error: {process.stderr}")239 240    @property241    def authorization(self) -> str:242        if self.beam_client_id:243            credential_str = self.beam_client_id + ":" + self.beam_client_secret244        else:245            credential_str = self.beam_client_secret246        return base64.b64encode(credential_str.encode()).decode()247 248    def _call(249        self,250        prompt: str,251        stop: Optional[list] = None,252        run_manager: Optional[CallbackManagerForLLMRun] = None,253        **kwargs: Any,254    ) -> str:255        """Call to Beam."""256        url = "https://apps.beam.cloud/" + self.app_id if self.app_id else self.url257        payload = {"prompt": prompt, "max_length": self.max_length}258        payload.update(kwargs)259        headers = {260            "Accept": "*/*",261            "Accept-Encoding": "gzip, deflate",262            "Authorization": "Basic " + self.authorization,263            "Connection": "keep-alive",264            "Content-Type": "application/json",265        }266 267        for _ in range(DEFAULT_NUM_TRIES):268            request = requests.post(url, headers=headers, data=json.dumps(payload))269            if request.status_code == 200:270                return request.json()["text"]271            time.sleep(DEFAULT_SLEEP_TIME)272        logger.warning("Unable to successfully call model.")273        return ""274