Update app.py

app.py

@@ -13,8 +13,7 @@ from botocore.exceptions import NoCredentialsError
 from functools import cached_property
 import base64
 from optimum.onnxruntime import ORTModelForCausalLM
-from optimum.bettertransformer import BetterTransformer
-
+import bitsandbytes as bnb
 
 AWS_ACCESS_KEY_ID = os.getenv("AWS_ACCESS_KEY_ID")
 AWS_SECRET_ACCESS_KEY = os.getenv("AWS_SECRET_ACCESS_KEY")
@@ -41,8 +40,8 @@ class GenerateRequest(BaseModel):
     num_return_sequences: int = 1
     do_sample: bool = False
     stop_sequences: list[str] = []
+    quantize: bool = True
     use_onnx: bool = False
-    use_bettertransformer: bool = True
     @field_validator("model_name")
     def model_name_cannot_be_empty(cls, v):
         if not v:
@@ -61,16 +60,19 @@ class S3ModelLoader:
         self.model_cache = {}
     def _get_s3_uri(self, model_name):
          return f"s3://{self.bucket_name}/{model_name.replace('/', '-')}"
-    async def _load_model_and_tokenizer(self, model_name, use_onnx, use_bettertransformer):
+    async def _load_model_and_tokenizer(self, model_name, quantize, use_onnx):
        s3_uri = self._get_s3_uri(model_name)
        try:
             config = AutoConfig.from_pretrained(s3_uri, local_files_only=False)
             if use_onnx:
                 model = ORTModelForCausalLM.from_pretrained(s3_uri, config=config, local_files_only=False).to(self.device)
+            elif quantize:
+                model = AutoModelForCausalLM.from_pretrained(
+                    s3_uri, config=config, local_files_only=False,
+                    load_in_8bit=True
+                    ).to(self.device)
             else:
                 model = AutoModelForCausalLM.from_pretrained(s3_uri, config=config, local_files_only=False).to(self.device)
-            if use_bettertransformer:
-                model = BetterTransformer.transform(model)
             tokenizer = AutoTokenizer.from_pretrained(s3_uri, config=config, local_files_only=False)
             if tokenizer.eos_token_id is not None and tokenizer.pad_token_id is None:
                 tokenizer.pad_token_id = config.pad_token_id or tokenizer.eos_token_id
@@ -81,10 +83,13 @@ class S3ModelLoader:
                 tokenizer = AutoTokenizer.from_pretrained(model_name, config=config, token=HUGGINGFACE_HUB_TOKEN)
                 if use_onnx:
                     model = ORTModelForCausalLM.from_pretrained(model_name, config=config, token=HUGGINGFACE_HUB_TOKEN).to(self.device)
+                elif quantize:
+                    model = AutoModelForCausalLM.from_pretrained(
+                        model_name, config=config, token=HUGGINGFACE_HUB_TOKEN,
+                        load_in_8bit=True
+                    ).to(self.device)
                 else:
                      model = AutoModelForCausalLM.from_pretrained(model_name, config=config, token=HUGGINGFACE_HUB_TOKEN).to(self.device)
-                if use_bettertransformer:
-                    model = BetterTransformer.transform(model)
                 if tokenizer.eos_token_id is not None and tokenizer.pad_token_id is None:
                     tokenizer.pad_token_id = config.pad_token_id or tokenizer.eos_token_id
                 return model, tokenizer
@@ -93,10 +98,10 @@ class S3ModelLoader:
     @cached_property
     def device(self):
         return torch.device("cpu")
-    async def get_model_and_tokenizer(self, model_name, use_onnx, use_bettertransformer):
-        key = f"{model_name}-{use_onnx}-{use_bettertransformer}"
+    async def get_model_and_tokenizer(self, model_name, quantize, use_onnx):
+        key = f"{model_name}-{quantize}-{use_onnx}"
         if key not in self.model_cache:
-            model, tokenizer = await self._load_model_and_tokenizer(model_name, use_onnx, use_bettertransformer)
+            model, tokenizer = await self._load_model_and_tokenizer(model_name, quantize, use_onnx)
             self.model_cache[key] = {"model":model, "tokenizer":tokenizer}
         return self.model_cache[key]["model"], self.model_cache[key]["tokenizer"]
     async def get_pipeline(self, model_name, task_type):
@@ -122,9 +127,9 @@ async def generate(request: GenerateRequest):
         num_return_sequences = request.num_return_sequences
         do_sample = request.do_sample
         stop_sequences = request.stop_sequences
+        quantize = request.quantize
         use_onnx = request.use_onnx
-        use_bettertransformer = request.use_bettertransformer
-        model, tokenizer = await model_loader.get_model_and_tokenizer(model_name, use_onnx, use_bettertransformer)
+        model, tokenizer = await model_loader.get_model_and_tokenizer(model_name, quantize, use_onnx)
         if "text-to-text" == task_type:
             generation_config = GenerationConfig(temperature=temperature,max_new_tokens=max_new_tokens,top_p=top_p,top_k=top_k,repetition_penalty=repetition_penalty,do_sample=do_sample,num_return_sequences=num_return_sequences,eos_token_id = tokenizer.eos_token_id)
             if stream: