Update app.py

app.py

@@ -1,34 +1,25 @@
 from fastapi import FastAPI, HTTPException
 from pydantic import BaseModel
 from typing import Optional
-from transformers import AutoTokenizer, AutoModelForCausalLM
 import torch
+from transformers import AutoTokenizer, AutoModelForCausalLM, AutoConfig
 
 app = FastAPI()
 
-# โหลดโมเดลและ tokenizer
 model_name = "scb10x/llama-3-typhoon-v1.5-8b-instruct"
 tokenizer = AutoTokenizer.from_pretrained(model_name)
 
-# ตรวจสอบว่ามี GPU หรือไม่
-device = "cuda" if torch.cuda.is_available() else "cpu"
+# ใช้ BitsAndBytes สำหรับ quantization
+config = AutoConfig.from_pretrained(model_name)
+config.quantization_config = BitsAndBytesConfig(load_in_8bit=True)
 
-# โหลดโมเดลด้วยการตั้งค่าที่เหมาะสม
-if device == "cuda":
-    model = AutoModelForCausalLM.from_pretrained(
-        model_name, 
-        torch_dtype=torch.float16, 
-        device_map="auto",
-        low_cpu_mem_usage=True
-    )
-else:
-    model = AutoModelForCausalLM.from_pretrained(
-        model_name,
-        torch_dtype=torch.float32,
-        low_cpu_mem_usage=True
-    )
-
-model.to(device)
+# โหลดโมเดลด้วย 8-bit quantization
+model = AutoModelForCausalLM.from_pretrained(
+    model_name,
+    config=config,
+    device_map="auto",
+    torch_dtype=torch.float16,
+)
 
 class Query(BaseModel):
     queryResult: Optional[dict] = None
@@ -42,15 +33,13 @@ async def webhook(query: Query):
         if not user_query:
             raise HTTPException(status_code=400, detail="No query text provided")
         
-        # สร้าง prompt และ generate ข้อความ
         prompt = f"Human: {user_query}\nAI:"
-        input_ids = tokenizer.encode(prompt, return_tensors="pt").to(device)
+        input_ids = tokenizer.encode(prompt, return_tensors="pt").to(model.device)
         
         with torch.no_grad():
             output = model.generate(input_ids, max_new_tokens=100, temperature=0.7)
         response = tokenizer.decode(output[0], skip_special_tokens=True)
         
-        # แยกส่วนที่เป็นคำตอบของ AI
         ai_response = response.split("AI:")[-1].strip()
         
         return {"fulfillmentText": ai_response}