add option for 8bit

app/api.py

@@ -136,13 +136,13 @@ async def text_generate(
     return {"generated_text": generated_text, "processing_time": time_diff}
 
 
-def get_text_generator(model_name: str, device: str = "cpu"):
+def get_text_generator(model_name: str, load_in_8bit: bool = False, device: str = "cpu"):
     hf_auth_token = os.getenv("HF_AUTH_TOKEN", False)
     print(f"hf_auth_token: {hf_auth_token}")
     print(f"Loading model with device: {device}...")
     tokenizer = AutoTokenizer.from_pretrained(model_name, use_auth_token=hf_auth_token)
     model = AutoModelForCausalLM.from_pretrained(model_name, pad_token_id=tokenizer.eos_token_id,
-                                            load_in_8bit=True, device_map="auto", use_auth_token=hf_auth_token)
+                                            load_in_8bit=load_in_8bit, device_map="auto", use_auth_token=hf_auth_token)
     # model.to(device)
     print("Model loaded")
     return model, tokenizer
@@ -156,7 +156,9 @@ config = get_config()
 device = "cuda" if torch.cuda.is_available() else "cpu"
 text_generator = {}
 for model_name in config["text-generator"]:
-    model, tokenizer = get_text_generator(model_name=config["text-generator"][model_name], device=device)
+    model, tokenizer = get_text_generator(model_name=config["text-generator"][model_name]["name"],
+                                          load_in_8bit=config["text-generator"][model_name]["load_in_8bit"],
+                                          device=device)
     text_generator[model_name] = {
         "model": model,
         "tokenizer": tokenizer

app/config.json

@@ -1,6 +1,12 @@
 {
   "text-generator": {
-    "indochat-tiny": "cahya/indochat-tiny",
-    "bloomz-1b1-instruct": "cahya/bloomz-1b7-instruct"
+    "indochat-tiny": {
+      "name": "cahya/indochat-tiny",
+      "load_in_8bit": false
+    },
+    "bloomz-1b1-instruct": {
+      "name": "cahya/bloomz-1b7-instruct",
+      "load_in_8bit": true
+    }
   }
 }