handler.py

from typing import Any
from typing import Dict
from typing import List


def report_gpu_usage() -> str:
    import os
    return os.popen("nvidia-smi").read()


class EndpointHandler:
    def __init__(self, path=""):
        import torch
        from transformers import AutoModelForSeq2SeqLM
        from transformers import AutoTokenizer
        import os

        os.environ["TOKENIZERS_PARALLELISM"] = "false"
        model_kwargs: dict[str, any] = dict()
        if torch.cuda.is_available():
            model_kwargs["load_in_8bit"] = True
            model_kwargs["torch_dtype"] = torch.bfloat16
            model_kwargs["device_map"] = "auto"
            model_kwargs["low_cpu_mem_usage"] = True
        self.model = AutoModelForSeq2SeqLM.from_pretrained(path, **model_kwargs)
        self.tokenizer = AutoTokenizer.from_pretrained(path)
        self.device = "cuda" if torch.cuda.is_available() else "cpu"
        print(f"Loaded model {path} to {self.device}")

    def __call__(self, data: Dict[str, Any]) -> List[Dict[str, any]]:
        import os
        os.environ["TOKENIZERS_PARALLELISM"] = "false"

        inputs = data.pop("inputs", data)
        input_ids = self.tokenizer(inputs, return_tensors="pt").input_ids.to(self.device)

        parameters = data.pop("parameters", None)
        if parameters is None:
            parameters = dict()

        with BlockTimer() as timer:
            outputs = self.model.generate(input_ids, **parameters)
        print(f"Inference elapsed: {round(timer.duration, 2)}")
        return_value: List[Dict[str, any]] = list()
        # postprocess the prediction
        gpu_info = report_gpu_usage()
        for output in outputs:
            prediction = self.tokenizer.decode(output, skip_special_tokens=True)
            entry = {"generated_text": prediction, "elapsed": timer.duration}
            if gpu_info:
                entry["gpu_info"] = gpu_info
                gpu_info = None
            return_value.append(entry)
        return return_value


class BlockTimer(object):
    def __enter__(self):
        import time
        self.start = time.perf_counter()
        return self

    def __exit__(self, typ, value, traceback):
        import time
        self.duration = time.perf_counter() - self.start


def _force_not_available() -> bool:
    return False


def test() -> None:
    import textwrap
    # torch.cuda.is_available = _force_not_available
    handler = EndpointHandler(path="bigscience/mt0-xl")
    parameters: dict[str, any] = {"max_length": 256, "min_length": 1,  #
                                  "no_repeat_ngram_size": 3,  #
                                  # "encoder_no_repeat_ngram_size": 7,  #
                                  "repetition_penalty": 3.5,  #
                                  # "num_beams": 1,  #
                                  # "top_p": 0.7,  # 0.3, 0.7
                                  "do_sample": True,
                                  "temperature": 0.1,
                                  "early_stopping": True, }  # parameters for text generation
    payload = {"inputs": f"{wall_of_text()}", "parameters": parameters}
    results = handler.__call__(payload)
    for entry in results[0].items():
        print()
        print(f"=== {entry[0]}")
        if entry[0] == "gpu_info":
            gpu_info_lines = entry[1].split("\n")
            for line in gpu_info_lines:
                if "Default |" in line:
                    print(line)
        else:
            print(textwrap.fill(str(
                    entry[1]), 140, drop_whitespace=False, replace_whitespace=False))


def wall_of_text() -> str:
    return """
    Write a journal article headline for the following.

    The present invention relates to compositions and methods for the treatment of the 
    Charcot-Marie-Tooth disease and related disorders. Charcot-Marie-Tooth disease (“CMT 
    Mining 
    of publicly available data, describing molecular mechanisms and pathological 
    manifestations 
    of the CMT1A disease, allowed us to prioritize a few functional cellular 
    modules-transcriptional regulation of PMP22 gene, PMP22 protein folding/degradation, 
    Schwann cell proliferation and apoptosis, death of neurons, extra-cellular matrix 
    deposition 
    and remodelling, immune response-as potential legitimate targets for CMT-relevant 
    therapeutic interventions. The combined impact of these deregulated functional modules on 
    onset and progression of pathological manifestations of Charcot-Marie-Tooth justifies a 
    potential efficacy of combinatorial CMT treatment. International patent application No. 
    PCT/EP2008/066457 describes a method of identifying drug candidates for the treatment of 
    the 
    Charcot-Marie-Tooth disease by building a dynamic model of the pathology and targeting 
    functional cellular pathways which are relevant in the regulation of CMT disease. 
    International patent application No. PCT/EP2008/066468 describes compositions for the 
    treatment of the Charcot-Marie-Tooth disease which comprise at least two compounds 
    selected 
    from the group of multiple drug candidates. The purpose of the present invention is to 
    provide new therapeutic combinations for treating CMT and related disorders. The invention 
    thus relates to compositions and methods for treating CMT and related disorders, 
    in particular toxic or traumatic neuropathy and amyotrophic lateral sclerosis, 
    using particular drug combinations. An object of this invention more specifically 
    relates to 
    a composition comprising baclofen, sorbitol and a compound selected from pilocarpine, 
    methimazole, mifepristone, naltrexone, rapamycin, flurbiprofen and ketoprofen, salts or 
    prodrugs thereof, for simultaneous, separate or sequential administration to a mammalian 
    subject. A particular object of the present invention relates to a composition comprising 
    baclofen, sorbitol and naltrexone, for simultaneous, separate or sequential administration 
    to a mammalian subject. Another object of the invention relates to a composition 
    comprising 
    (a) rapamycin, (b) mifepristone or naltrexone, and © a PMP22 modulator, for simultaneous, 
    separate or sequential administration to a mammalian subject. In a particular embodiment, 
    the PMP22 modulator is selected from acetazolamide, albuterol, amiloride, 
    aminoglutethimide, 
    amiodarone, aztreonam, baclofen, balsalazide, betaine, bethanechol, bicalutamide, 
    bromocriptine, bumetanide, buspirone, carbachol, carbamazepine, carbimazole, cevimeline, 
    ciprofloxacin, clonidine, curcumin, cyclosporine A, diazepam, diclofenac, dinoprostone, 
    disulfiram, D-sorbitol, dutasteride, estradiol, exemestane, felbamate, fenofibrate, 
    finasteride, flumazenil, flunitrazepam, flurbiprofen, furosemide, gabapentingabapentin, 
    galantamine, haloperidol, ibuprofen, isoproterenol, ketoconazole, ketoprofen, L-carnitine, 
    liothyronine (T3), lithium, losartan, loxapine, meloxicam, metaproterenol, metaraminol, 
    metformin, methacholine, methimazole, methylergonovine, metoprolol, metyrapone, 
    miconazole, 
    mifepristone, nadolol, naloxone, naltrexone; norfloxacin, pentazocine, phenoxybenzamine, 
    phenylbutyrate, pilocarpine, pioglitazone, prazosin, propylthiouracil, raloxifene, 
    rapamycin, rifampin, simvastatin, spironolactone, tacrolimus, tamoxifen, trehalose, 
    trilostane, valproic acid, salts or prodrugs thereof. 1. A method of improving nerve 
    regeneration in a human subject suffering from amyotrophic lateral sclerosis, 
    or a neuropathy selected from an idiopathic neuropathy, diabetic neuropathy, 
    a toxic neuropathy, a neuropathy induced by a drug treatment, a neuropathy provoked by 
    HIV, 
    a neuropathy provoked by radiation, a neuropathy provoked by heavy metals, a neuropathy 
    provoked by vitamin deficiency states, or a traumatic neuropathy, comprising administering 
    to the human subject an amount of a composition effective to improve nerve regeneration; 
    and 
    wherein the composition comprises baclofen or a pharmaceutically acceptable salt thereof 
    in 
    an amount from 1 to 300 mg/kg of the human subject per day; D-sorbitol or a 
    pharmaceutically 
    acceptable salt thereof; and naltrexone or a pharmaceutically acceptable salt thereof in 
    an 
    amount from 1 to 100 mg/kg of the human subject per day. 2. The method of claim 1, 
    wherein the composition further comprises a pharmaceutically suitable excipient or 
    carrier. 
    3. The method of claim 2, wherein the composition is formulated with a drug eluting 
    polymer, 
    a biomolecule, a micelle or liposome-forming lipids or oil in water emulsions, 
    or pegylated 
    or solid nanoparticles or microparticles for oral or parenteral or intrathecal 
    administration. 4. The method of claim 1, wherein the subject suffers from a traumatic 
    neuropathy arising from brain injury, spinal cord injury, or an injury to peripheral 
    nerves. 
    5. The method of claim 1, wherein the D-sorbitol or a pharmaceutically acceptable salt 
    thereof is D-sorbitol. 6. The method of claim 1, wherein the composition is formulated for 
    oral administration. 7. The method of claim 6, wherein the composition is a liquid 
    formulation. 8. The method of claim 1, wherein baclofen or a pharmaceutically acceptable 
    salt thereof, D-sorbitol or a pharmaceutically acceptable salt thereof, and naltrexone 
    or a 
    pharmaceutically acceptable salt thereof are the sole active ingredients. 9. The method of 
    claim 1, comprising administering to the human subject baclofen or a pharmaceutically 
    acceptable salt thereof in an amount from 10 to 200 mg/kg of the human subject per day and 
    naltrexone or a pharmaceutically acceptable salt thereof in an amount from 1 to 50 mg/kg 
    of 
    the human subject per day. 10. The method of claim 1, comprising administering to the 
    human 
    subject baclofen or a pharmaceutically acceptable salt thereof in an amount from 10 to 200 
    mg/kg of the human subject per day and naltrexone or a pharmaceutically acceptable salt 
    thereof in an amount from 1 to 50 mg/kg of the human subject per day. 11. The method of 
    claim 1, comprising administering to the human subject baclofen or a pharmaceutically 
    acceptable salt thereof in an amount from 60 mg to 18 mg per day and naltrexone or a 
    pharmaceutically acceptable salt thereof in an amount from 60 mg to 6 mg per day. 12. The 
    method of claim 1, comprising administering to the human subject baclofen or a 
    pharmaceutically acceptable salt thereof in an amount from 60 mg to 12 mg per day and 
    naltrexone or a pharmaceutically acceptable salt thereof in an amount from 60 mg to 3 mg 
    per 
    day. 13. The method of claim 10, wherein baclofen or a pharmaceutically acceptable salt 
    thereof, D-sorbitol or a pharmaceutically acceptable salt thereof, and naltrexone or a 
    pharmaceutically acceptable salt thereof are administered orally to the human subject. 14. 
    The method of claim 10, wherein baclofen or a pharmaceutically acceptable salt thereof, 
    D-sorbitol or a pharmaceutically acceptable salt thereof, and naltrexone or a 
    pharmaceutically acceptable salt thereof are administered separately to the human subject. 
    15. The method of claim 13, wherein baclofen or a pharmaceutically acceptable salt 
    thereof, 
    D-sorbitol or a pharmaceutically acceptable salt thereof, and naltrexone or a 
    pharmaceutically acceptable salt thereof are formulated in a liquid formulation. 16. The 
    method of claim 15, wherein baclofen or a pharmaceutically acceptable salt thereof, 
    D-sorbitol or a pharmaceutically acceptable salt thereof, and naltrexone or a 
    pharmaceutically acceptable salt thereof are administered to the human subject in divided 
    doses. 17. The method of claim 15, wherein baclofen or a pharmaceutically acceptable salt 
    thereof, D-sorbitol or a pharmaceutically acceptable salt thereof, and naltrexone or a 
    pharmaceutically acceptable salt thereof are administered to the human subject in divided 
    doses two times daily.
""".replace("\n", " ")


if __name__ == '__main__':
    test()