Update app.py

app.py

@@ -14,7 +14,7 @@ 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,7 +41,6 @@ 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")
@@ -62,17 +61,12 @@ 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, quantize, use_onnx, use_bettertransformer):
+    async def _load_model_and_tokenizer(self, model_name, use_onnx, use_bettertransformer):
        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:
@@ -87,11 +81,6 @@ 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:
@@ -104,10 +93,10 @@ class S3ModelLoader:
     @cached_property
     def device(self):
         return torch.device("cpu")
-    async def get_model_and_tokenizer(self, model_name, quantize, use_onnx, use_bettertransformer):
-        key = f"{model_name}-{quantize}-{use_onnx}-{use_bettertransformer}"
+    async def get_model_and_tokenizer(self, model_name, use_onnx, use_bettertransformer):
+        key = f"{model_name}-{use_onnx}-{use_bettertransformer}"
         if key not in self.model_cache:
-            model, tokenizer = await self._load_model_and_tokenizer(model_name, quantize, use_onnx, use_bettertransformer)
+            model, tokenizer = await self._load_model_and_tokenizer(model_name, use_onnx, use_bettertransformer)
             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):
@@ -133,10 +122,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, quantize, use_onnx, use_bettertransformer)
+        model, tokenizer = await model_loader.get_model_and_tokenizer(model_name, use_onnx, use_bettertransformer)
         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: