Update app.py

app.py

@@ -2,11 +2,9 @@ from fastapi import FastAPI, HTTPException
 from pydantic import BaseModel
 from typing import Optional
 import torch
-from transformers import AutoTokenizer, AutoModelForCausalLM, AutoConfig
+from transformers import AutoTokenizer, AutoModelForCausalLM, BitsAndBytesConfig
 import logging
-import os
 
-# ตั้งค่า logging
 logging.basicConfig(level=logging.INFO)
 logger = logging.getLogger(__name__)
 
@@ -16,32 +14,21 @@ try:
     model_name = "scb10x/llama-3-typhoon-v1.5-8b-instruct"
     tokenizer = AutoTokenizer.from_pretrained(model_name)
 
-    # ตรวจสอบว่ามี GPU หรือไม่
-    if torch.cuda.is_available():
-        logger.info("GPU is available. Using CUDA.")
-        device = "cuda"
-    else:
-        logger.info("No GPU found. Using CPU.")
-        device = "cpu"
+    device = "cuda" if torch.cuda.is_available() else "cpu"
+    logger.info(f"Using device: {device}")
 
-    # กำหนดการตั้งค่าสำหรับการโหลดโมเดล
-    model_kwargs = {
-        "torch_dtype": torch.float32 if device == "cpu" else torch.float16,
-        "low_cpu_mem_usage": True,
-    }
-
-    if device == "cuda":
-        from transformers import BitsAndBytesConfig
-        model_kwargs["quantization_config"] = BitsAndBytesConfig(load_in_8bit=True)
+    # 4-bit quantization configuration
+    quantization_config = BitsAndBytesConfig(
+        load_in_4bit=True,
+        bnb_4bit_compute_dtype=torch.float16
+    )
 
-    # โหลดโมเดล
     model = AutoModelForCausalLM.from_pretrained(
         model_name,
-        device_map="auto" if device == "cuda" else None,
-        **model_kwargs
+        quantization_config=quantization_config,
+        device_map="auto",
+        low_cpu_mem_usage=True,
     )
-
-    model.to(device)
     logger.info(f"Model loaded successfully on {device}")
 except Exception as e:
     logger.error(f"Error loading model: {str(e)}")