Update main.py

main.py

@@ -1,93 +1,110 @@
-from fastapi import FastAPI
-from fastapi.responses import JSONResponse
-import firebase_admin
-from firebase_admin import credentials, firestore
-from transformers import pipeline
-from pydantic import BaseModel
-import os
-from huggingface_hub import login
-from transformers import AutoTokenizer, AutoModelForSeq2SeqLM   # AutoModelForCausalLM,
+# from fastapi import FastAPI
+# from fastapi.responses import JSONResponse
+# import firebase_admin
+# from firebase_admin import credentials, firestore
+# from transformers import pipeline
+# from pydantic import BaseModel
+# import os
+# from huggingface_hub import login
+# from transformers import AutoTokenizer, AutoModelForSeq2SeqLM   # AutoModelForCausalLM,
 
-# Load Firebase
-cred = credentials.Certificate("firebase_config.json")
-firebase_admin.initialize_app(cred)
-db = firestore.client()
+# # Load Firebase
+# cred = credentials.Certificate("firebase_config.json")
+# firebase_admin.initialize_app(cred)
+# db = firestore.client()
 
-# Đăng nhập vào Hugging Face (nếu cần)
-HF_TOKEN = os.getenv("HF_TOKEN")
-if HF_TOKEN:
-    login(HF_TOKEN)
-
-# Load AI Model
-tokenizer = AutoTokenizer.from_pretrained("google/flan-t5-base",token=HF_TOKEN)
-ai_model = AutoModelForSeq2SeqLM.from_pretrained("google/flan-t5-base",token=HF_TOKEN)
-
-app = FastAPI()
-
-# Request Models
-class UpdateBiosRequest(BaseModel):
-    user_id: str
-    bios: str
-
-class FocusHistoryRequest(BaseModel):
-    user_id: str
-    time_start: str
-    total_time: int
-
-class AIRequest(BaseModel):
-    user_id: str
-
-# API cập nhật bios
-@app.post("/update_bios")
-async def update_bios(request: UpdateBiosRequest):
-    user_ref = db.collection("user_profiles").document(request.user_id)
-    user_ref.set({"bios": request.bios}, merge=True)
-    return {"message": "Cập nhật bios thành công"}
-
-# API thêm lịch sử focus
-@app.post("/add_focus_history")
-async def add_focus_history(request: FocusHistoryRequest):
-    user_ref = db.collection("user_profiles").document(request.user_id)
-    user_doc = user_ref.get()
-    data = user_doc.to_dict() or {}
-    focus_history = data.get("focus_history", [])
-    focus_history.append({"time_start": request.time_start, "total_time": request.total_time})
-    user_ref.set({"focus_history": focus_history}, merge=True)
-    return {"message": "Thêm lịch sử focus thành công"}
-
-# API lấy dữ liệu người dùng
-@app.get("/get_user_data")
-async def get_user_data(user_id: str):
-    user_doc = db.collection("user_profiles").document(user_id).get()
-    data = user_doc.to_dict() or {}
-    return {
-        "bios": data.get("bios", "Chưa có bios."),
-        "focus_history": data.get("focus_history", [])
-    }
-
-# API AI tư vấn
-@app.post("/ai_personal_advice")
-async def ai_personal_advice(request: AIRequest):
-    user_doc = db.collection("user_profiles").document(request.user_id).get()
-    data = user_doc.to_dict() or {}
-    bios = data.get("bios", "Chưa có bios.")
-    focus_history = data.get("focus_history", [])
-    focus_text = "\n".join([f"- {f['time_start']}: {f['total_time']} phút" for f in focus_history])
+# # Đăng nhập vào Hugging Face (nếu cần)
+# HF_TOKEN = os.getenv("HF_TOKEN")
+# if HF_TOKEN:
+#     login(HF_TOKEN)
+
+# # Load AI Model
+# tokenizer = AutoTokenizer.from_pretrained("google/flan-t5-base",token=HF_TOKEN)
+# ai_model = AutoModelForSeq2SeqLM.from_pretrained("google/flan-t5-base",token=HF_TOKEN)
+
+# app = FastAPI()
+
+# # Request Models
+# class UpdateBiosRequest(BaseModel):
+#     user_id: str
+#     bios: str
+
+# class FocusHistoryRequest(BaseModel):
+#     user_id: str
+#     time_start: str
+#     total_time: int
+
+# class AIRequest(BaseModel):
+#     user_id: str
+
+# # API cập nhật bios
+# @app.post("/update_bios")
+# async def update_bios(request: UpdateBiosRequest):
+#     user_ref = db.collection("user_profiles").document(request.user_id)
+#     user_ref.set({"bios": request.bios}, merge=True)
+#     return {"message": "Cập nhật bios thành công"}
+
+# # API thêm lịch sử focus
+# @app.post("/add_focus_history")
+# async def add_focus_history(request: FocusHistoryRequest):
+#     user_ref = db.collection("user_profiles").document(request.user_id)
+#     user_doc = user_ref.get()
+#     data = user_doc.to_dict() or {}
+#     focus_history = data.get("focus_history", [])
+#     focus_history.append({"time_start": request.time_start, "total_time": request.total_time})
+#     user_ref.set({"focus_history": focus_history}, merge=True)
+#     return {"message": "Thêm lịch sử focus thành công"}
+
+# # API lấy dữ liệu người dùng
+# @app.get("/get_user_data")
+# async def get_user_data(user_id: str):
+#     user_doc = db.collection("user_profiles").document(user_id).get()
+#     data = user_doc.to_dict() or {}
+#     return {
+#         "bios": data.get("bios", "Chưa có bios."),
+#         "focus_history": data.get("focus_history", [])
+#     }
+
+# # API AI tư vấn
+# @app.post("/ai_personal_advice")
+# async def ai_personal_advice(request: AIRequest):
+#     user_doc = db.collection("user_profiles").document(request.user_id).get()
+#     data = user_doc.to_dict() or {}
+#     bios = data.get("bios", "Chưa có bios.")
+#     focus_history = data.get("focus_history", [])
+#     focus_text = "\n".join([f"- {f['time_start']}: {f['total_time']} phút" for f in focus_history])
     
-    prompt = f"""
-    Thông tin người dùng:
-    - Bios: {bios}
-    - Lịch sử focus:
-    {focus_text}
+#     prompt = f"""
+#     Thông tin người dùng:
+#     - Bios: {bios}
+#     - Lịch sử focus:
+#     {focus_text}
     
-    Hãy tư vấn cách cải thiện hiệu suất làm việc dựa trên thông tin trên.
-    """
-    input_ids = tokenizer(combined_information, return_tensors="pt")
-    response = ai_model.generate(**input_ids, max_new_tokens=500)
-    # response = ai_model(prompt, max_length=200)
-    return {"advice": tokenizer.decode(response[0])}
-
-# Trang chủ
-@app.get("/")
-async def home():
-    return JSONResponse(content={"message": "Welcome to the Recommendation API!"})
+#     Hãy tư vấn cách cải thiện hiệu suất làm việc dựa trên thông tin trên.
+#     """
+#     input_ids = tokenizer(combined_information, return_tensors="pt")
+#     response = ai_model.generate(**input_ids, max_new_tokens=500)
+#     # response = ai_model(prompt, max_length=200)
+#     return {"advice": tokenizer.decode(response[0])}
+
+# # Trang chủ
+# @app.get("/")
+# async def home():
+#     return JSONResponse(content={"message": "Welcome to the Recommendation API!"})
+
+
+
+from transformers import AutoTokenizer, AutoModelForSeq2SeqLM
+
+# Kiểm tra biến môi trường HF_TOKEN
+import os
+HF_TOKEN = os.getenv("HF_TOKEN")
+print("HF_TOKEN:", HF_TOKEN is not None)
+
+# Load model
+model_name = "google/flan-t5-base"
+tokenizer = AutoTokenizer.from_pretrained(model_name, token=HF_TOKEN)
+ai_model = AutoModelForSeq2SeqLM.from_pretrained(model_name, token=HF_TOKEN)
+
+print("Model loaded successfully!")
+