Update app.py

app.py

@@ -12,8 +12,6 @@ from huggingface_hub import login
 from botocore.exceptions import NoCredentialsError
 from functools import cached_property
 import base64
-from optimum.onnxruntime import ORTModelForCausalLM
-import bitsandbytes as bnb
 
 AWS_ACCESS_KEY_ID = os.getenv("AWS_ACCESS_KEY_ID")
 AWS_SECRET_ACCESS_KEY = os.getenv("AWS_SECRET_ACCESS_KEY")
@@ -40,7 +38,7 @@ class GenerateRequest(BaseModel):
     num_return_sequences: int = 1
     do_sample: bool = False
     stop_sequences: list[str] = []
-    quantize: bool = True
+    quantize: bool = False
     use_onnx: bool = False
     @field_validator("model_name")
     def model_name_cannot_be_empty(cls, v):
@@ -64,15 +62,7 @@ class S3ModelLoader:
        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)
+            model = AutoModelForCausalLM.from_pretrained(s3_uri, config=config, local_files_only=False).to(self.device)
             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,15 +71,7 @@ class S3ModelLoader:
             try:
                 config = AutoConfig.from_pretrained(model_name, token=HUGGINGFACE_HUB_TOKEN)
                 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)
+                model = AutoModelForCausalLM.from_pretrained(model_name, config=config, token=HUGGINGFACE_HUB_TOKEN).to(self.device)
                 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