Update app.py

app.py

@@ -3,8 +3,8 @@ from pydantic import BaseModel
 from typing import Optional
 import torch
 from transformers import AutoTokenizer, AutoModelForCausalLM, AutoConfig
-from transformers import BitsAndBytesConfig  # เพิ่มการ import นี้
 import logging
+import os
 
 # ตั้งค่า logging
 logging.basicConfig(level=logging.INFO)
@@ -16,18 +16,33 @@ try:
     model_name = "scb10x/llama-3-typhoon-v1.5-8b-instruct"
     tokenizer = AutoTokenizer.from_pretrained(model_name)
 
-    # ใช้ BitsAndBytes สำหรับ quantization
-    config = AutoConfig.from_pretrained(model_name)
-    config.quantization_config = BitsAndBytesConfig(load_in_8bit=True)
+    # ตรวจสอบว่ามี 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"
 
-    # โหลดโมเดลด้วย 8-bit quantization
+    # กำหนดการตั้งค่าสำหรับการโหลดโมเดล
+    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)
+
+    # โหลดโมเดล
     model = AutoModelForCausalLM.from_pretrained(
         model_name,
-        config=config,
-        device_map="auto",
-        torch_dtype=torch.float16,
+        device_map="auto" if device == "cuda" else None,
+        **model_kwargs
     )
-    logger.info("Model loaded successfully")
+
+    model.to(device)
+    logger.info(f"Model loaded successfully on {device}")
 except Exception as e:
     logger.error(f"Error loading model: {str(e)}")
     raise
@@ -45,7 +60,7 @@ async def webhook(query: Query):
             raise HTTPException(status_code=400, detail="No query text provided")
         
         prompt = f"Human: {user_query}\nAI:"
-        input_ids = tokenizer.encode(prompt, return_tensors="pt").to(model.device)
+        input_ids = tokenizer.encode(prompt, return_tensors="pt").to(device)
         
         with torch.no_grad():
             output = model.generate(input_ids, max_new_tokens=100, temperature=0.7)