Update app.py

app.py

@@ -1,35 +1,36 @@
-from transformers import AutoModelForCausalLM, AutoTokenizer
 from fastapi import FastAPI
+from pydantic import BaseModel
+from transformers import AutoModelForCausalLM, AutoTokenizer
+import torch
+import os
+
 app = FastAPI()
 model_name = "Qwen/Qwen3-4B-Instruct-2507"
 
-# load the tokenizer and the model
+# Load tokenizer và model (CPU)
 tokenizer = AutoTokenizer.from_pretrained(model_name)
 model = AutoModelForCausalLM.from_pretrained(
     model_name,
-    dtype="auto",
-    device_map="auto"
+    device_map={"": "cpu"},  # Spaces Free không có GPU
+    torch_dtype=torch.float32
 )
 
-# prepare the model input
-prompt = "trả lời bằng tiếng việt, ngắn gọn."
-messages = [
-    {"role": "user", "content": prompt}
-]
-text = tokenizer.apply_chat_template(
-    messages,
-    tokenize=False,
-    add_generation_prompt=True,
-)
-model_inputs = tokenizer([text], return_tensors="pt").to(model.device)
+class ChatRequest(BaseModel):
+    message: str
 
-# conduct text completion
-generated_ids = model.generate(
-    **model_inputs,
-    max_new_tokens=200
-)
-output_ids = generated_ids[0][len(model_inputs.input_ids[0]):].tolist() 
+@app.get("/")
+def read_root():
+    return {"message": "Ứng dụng đang chạy!"}
+
+@app.post("/chat")
+async def chat(request: ChatRequest):
+    messages = [{"role": "user", "content": request.message}]
+    text = tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)
+    model_inputs = tokenizer([text], return_tensors="pt").to(model.device)
 
-content = tokenizer.decode(output_ids, skip_special_tokens=True)
+    with torch.inference_mode():
+        generated_ids = model.generate(**model_inputs, max_new_tokens=200)
 
-print("content:", content)
+    output_ids = generated_ids[0][len(model_inputs.input_ids[0]):].tolist()
+    content = tokenizer.decode(output_ids, skip_special_tokens=True)
+    return {"response": content}