app.py

from fastapi import FastAPI, HTTPException
from pydantic import BaseModel
from typing import Optional
from transformers import AutoTokenizer, AutoModelForCausalLM
import torch

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"

# โหลดโมเดลด้วยการตั้งค่าที่เหมาะสม
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)

class Query(BaseModel):
    queryResult: Optional[dict] = None
    queryText: Optional[str] = None

@app.post("/webhook")
async def webhook(query: Query):
    try:
        user_query = query.queryResult.get('queryText') if query.queryResult else query.queryText
        
        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)
        
        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}
    except Exception as e:
        raise HTTPException(status_code=500, detail=str(e))

if __name__ == "__main__":
    import uvicorn
    uvicorn.run(app, host="0.0.0.0", port=7860)