Update app.py

app.py

@@ -1,41 +1,87 @@
 from fastapi import FastAPI, Request
+from pydantic import BaseModel
 from transformers import AutoModelForCausalLM, AutoTokenizer
 import torch
-from pydantic import BaseModel
+from datetime import datetime
+import requests
+import pytz
+
 app = FastAPI()
 
 # --- Tải model ---
 MODEL_NAME = "Qwen/Qwen2.5-1.5B-Instruct"
+tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME, trust_remote_code=True)
+model = AutoModelForCausalLM.from_pretrained(
+    MODEL_NAME,
+    torch_dtype=torch.float16 if torch.cuda.is_available() else torch.float32,
+    device_map="auto"
+)
 
-try:
-    tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME, trust_remote_code=True)
-    model = AutoModelForCausalLM.from_pretrained(
-        MODEL_NAME,
-        torch_dtype=torch.float16 if torch.cuda.is_available() else torch.float32,
-        device_map="auto"
-    )
-except Exception as e:
-    raise RuntimeError(f"Lỗi khi tải model {MODEL_NAME}: {e}")
+# --- Request body ---
+class GenerateRequest(BaseModel):
+    prompt: str
+
+# --- Hàm lấy vị trí từ IP ---
+def get_location_from_ip(ip: str):
+    try:
+        res = requests.get(f"https://ipinfo.io/{ip}?token=2d478668dc5662").json()
+        city = res.get("city", "Ho Chi Minh City")
+        country = res.get("country", "VN")
+        timezone = res.get("timezone", "Asia/Ho_Chi_Minh")
+        return city, country, timezone
+    except:
+        return "Ho Chi Minh City", "VN", "Asia/Ho_Chi_Minh"
 
 @app.get("/")
 def home():
     return {"message": "API FastAI đã sẵn sàng phục vụ Túc chủ 👑"}
 
 @app.post("/generate")
-async def generate(request: Request):
-    data = await request.json()
-    prompt = data.get("prompt", "")
+async def generate(req: GenerateRequest, request: Request):
+    prompt = req.prompt.strip()
 
-    if not prompt:
-        return {"error": "Thiếu prompt!"}
+    # Lấy IP từ request
+    client_ip = request.client.host
+    city, country, timezone = get_location_from_ip(client_ip)
 
-# Thêm quy định vào prompt
-    final_prompt = f"Trả lời bằng tiếng Việt, ngắn gọn: {prompt}"
+    # --- Logic đặc biệt ---
+    # 1. Thời gian hiện tại
+    if "mấy giờ" in prompt or "thời gian" in prompt:
+        tz = pytz.timezone(timezone)
+        current_time = datetime.now(tz).strftime("%H:%M")
+        return {"response": f"Bây giờ ở {city} là {current_time}"}
 
+    # 2. Ngày hôm nay
+    if "ngày hôm nay" in prompt or "hôm nay là ngày" in prompt:
+        tz = pytz.timezone(timezone)
+        today = datetime.now(tz).strftime("%d/%m/%Y")
+        return {"response": f"Hôm nay ở {city} là ngày {today}"}
+
+    # 3. Thời tiết
+    if "thời tiết" in prompt:
+        try:
+            api_key = "YOUR_OPENWEATHER_API_KEY"
+            url = f"http://api.openweathermap.org/data/2.5/weather?q={city}&appid={b7d49ff5de091794a9adc8ea62ef0ac7}&units=metric&lang=vi"
+            res = requests.get(url).json()
+            temp = res["main"]["temp"]
+            desc = res["weather"][0]["description"]
+            return {"response": f"Thời tiết ở {city}: {desc}, nhiệt độ {temp}°C"}
+        except:
+            return {"response": "Không lấy được dữ liệu thời tiết."}
+
+    # --- Fallback sang mô hình ---
+    final_prompt = f"Trả lời bằng tiếng Việt, ngắn gọn: {prompt}"
     inputs = tokenizer(final_prompt, return_tensors="pt").to(model.device)
-    outputs = model.generate(**inputs,
-                             max_new_tokens=50,temperature=0.3,top_p=0.8,do_sample= True,repetition_penalty=1.2,eos_token_id=tokenizer.eos_token_id
-                            )
+    outputs = model.generate(
+        **inputs,
+        max_new_tokens=200,
+        temperature=0.3,
+        top_p=0.8,
+        do_sample=True,
+        repetition_penalty=1.2,
+        eos_token_id=tokenizer.eos_token_id
+    )
+
     result = tokenizer.decode(outputs[0], skip_special_tokens=True)
     clean_result = result[len(final_prompt):].strip()
     return {"response": clean_result}