Update app.py

app.py

@@ -3,6 +3,7 @@ from fastapi.middleware.cors import CORSMiddleware
 from pydantic import BaseModel
 from typing import Optional, List
 from transformers import AutoTokenizer, AutoModelForCausalLM
+from transformers import BitsAndBytesConfig
 from peft import PeftModel
 import torch
 import os
@@ -39,12 +40,15 @@ try:
     tokenizer = AutoTokenizer.from_pretrained(BASE_MODEL, use_fast=False)
     tokenizer.pad_token = tokenizer.eos_token
 
-    base_model = AutoModelForCausalLM.from_pretrained(
-        BASE_MODEL,
-        torch_dtype=torch.float16,
-        device_map="auto"
-    )
+quant_config = BitsAndBytesConfig(load_in_4bit=True)
 
+base_model = AutoModelForCausalLM.from_pretrained(
+    BASE_MODEL,
+    device_map="auto",
+    trust_remote_code=True,
+    offload_folder="offload",
+    quantization_config=quant_config
+    
     model = PeftModel.from_pretrained(base_model, LORA_MODEL)
     model.eval()