Update app.py

app.py

@@ -11,7 +11,10 @@ import docx2txt
 import chromadb
 from chromadb.config import Settings
 from shutil import rmtree
+import gc
+import torch
 
+# Optimization: Import only what's needed
 from langchain_google_genai import ChatGoogleGenerativeAI
 from langchain_chroma import Chroma
 from langchain_community.document_loaders import PyPDFLoader
@@ -24,213 +27,319 @@ from langchain_core.prompts import ChatPromptTemplate, MessagesPlaceholder
 from langchain_core.messages import HumanMessage, AIMessage
 from langchain_core.documents import Document
 from langchain_huggingface import HuggingFaceEmbeddings
-from langchain.retrievers import ContextualCompressionRetriever
-from langchain.retrievers.document_compressors import CrossEncoderReranker
-from langchain_community.cross_encoders import HuggingFaceCrossEncoder
+
+# Bỏ CrossEncoder để giảm memory, thay bằng các kỹ thuật khác
+# from langchain.retrievers import ContextualCompressionRetriever
+# from langchain.retrievers.document_compressors import CrossEncoderReranker
+# from langchain_community.cross_encoders import HuggingFaceCrossEncoder
 
 GOOGLE_API_KEY = os.getenv("GOOGLE_API_KEY")
 DATA_PATH = "medical_data"
 DB_PATH = "chroma_db"
-MAX_HISTORY_TURNS = 6
+MAX_HISTORY_TURNS = 5  # Giảm để tăng tốc
 FORCE_REBUILD_DB = False
 
 logging.basicConfig(level=logging.INFO, format="%(asctime)s [%(levelname)s] %(message)s")
 
 def process_excel_file(file_path: str, filename: str) -> list[Document]:
-    """
-    Xử lý Excel thông minh: Biến mỗi dòng thành một Document riêng biệt
-    giúp tìm kiếm chính xác từng bản ghi thuốc/bệnh nhân.
-    """
+    """Tối ưu hóa xử lý Excel để tăng độ chính xác"""
     docs = []
     try:
         if file_path.endswith(".csv"):
-            df = pd.read_csv(file_path)
+            df = pd.read_csv(file_path, encoding='utf-8')
         else:
             df = pd.read_excel(file_path)
-
-        df.dropna(how='all', inplace=True)
-        df.fillna("Không có thông tin", inplace=True)
-
-        for idx, row in df.iterrows():
-            content_parts = []
-            for col_name, val in row.items():
-                clean_val = str(val).strip()
-                if clean_val and clean_val.lower() != "nan":
-                    content_parts.append(f"{col_name}: {clean_val}")
-            
-            if content_parts:
-                page_content = f"Dữ liệu từ file {filename} (Dòng {idx+1}):\n" + "\n".join(content_parts)
-                metadata = {"source": filename, "row": idx+1, "type": "excel_record"}
-                docs.append(Document(page_content=page_content, metadata=metadata))
-                
+        
+        # Xử lý thông minh cho dữ liệu y tế
+        # Phát hiện loại dữ liệu tự động
+        if any(col in df.columns.str.lower() for col in ['tên thuốc', 'thuốc', 'drug']):
+            # Dữ liệu thuốc
+            for idx, row in df.iterrows():
+                content = f"THÔNG TIN THUỐC - Dòng {idx+1}:\n"
+                for col in df.columns:
+                    if pd.notna(row[col]):
+                        content += f"{col}: {str(row[col]).strip()}\n"
+                docs.append(Document(
+                    page_content=content,
+                    metadata={"source": filename, "row": idx+1, "type": "drug_info"}
+                ))
+        elif any(col in df.columns.str.lower() for col in ['bệnh nhân', 'patient', 'mã bn']):
+            # Dữ liệu bệnh nhân
+            for idx, row in df.iterrows():
+                content = f"HỒ SƠ BỆNH NHÂN - Dòng {idx+1}:\n"
+                for col in df.columns:
+                    if pd.notna(row[col]):
+                        content += f"{col}: {str(row[col]).strip()}\n"
+                docs.append(Document(
+                    page_content=content,
+                    metadata={"source": filename, "row": idx+1, "type": "patient_record"}
+                ))
+        else:
+            # Dữ liệu chung
+            for idx, row in df.iterrows():
+                content_parts = [f"{col}: {str(row[col]).strip()}" 
+                               for col in df.columns if pd.notna(row[col])]
+                if content_parts:
+                    docs.append(Document(
+                        page_content=f"Dữ liệu từ {filename} (Dòng {idx+1}):\n" + "\n".join(content_parts),
+                        metadata={"source": filename, "row": idx+1, "type": "general_data"}
+                    ))
+                    
     except Exception as e:
         logging.error(f"Lỗi xử lý Excel {filename}: {e}")
     
     return docs
 
 def load_documents_from_folder(folder_path: str) -> list[Document]:
-    logging.info(f"--- Bắt đầu quét thư mục: {folder_path} ---")
-    documents: list[Document] = []
+    """Tải và xử lý tài liệu với metadata phong phú"""
+    documents = []
     if not os.path.exists(folder_path):
         os.makedirs(folder_path, exist_ok=True)
         return []
-        
-    for root, _, files in os.walk(folder_path):
-        for filename in files:
-            file_path = os.path.join(root, filename)
-            filename_lower = filename.lower()
-            try:
-                if filename_lower.endswith(".pdf"):
-                    loader = PyPDFLoader(file_path)
-                    docs = loader.load()
-                    for d in docs: d.metadata["source"] = filename
-                    documents.extend(docs)
-                
-                elif filename_lower.endswith(".docx"):
-                    text = docx2txt.process(file_path)
-                    if text.strip(): 
-                        documents.append(Document(page_content=text, metadata={"source": filename}))
-                
-                elif filename_lower.endswith((".xlsx", ".xls", ".csv")):
-                    excel_docs = process_excel_file(file_path, filename)
-                    documents.extend(excel_docs)
-                    
-                elif filename_lower.endswith((".txt", ".md")):
-                    with open(file_path, "r", encoding="utf-8") as f: text = f.read()
-                    if text.strip(): 
-                        documents.append(Document(page_content=text, metadata={"source": filename}))
+    
+    # Ưu tiên xử lý theo thứ tự để tăng độ chính xác
+    file_extensions = ['.pdf', '.docx', '.xlsx', '.xls', '.csv', '.txt']
+    
+    for ext in file_extensions:
+        for root, _, files in os.walk(folder_path):
+            for filename in files:
+                if filename.lower().endswith(ext):
+                    file_path = os.path.join(root, filename)
+                    try:
+                        if filename.lower().endswith(".pdf"):
+                            loader = PyPDFLoader(file_path)
+                            docs = loader.load()
+                            for i, d in enumerate(docs):
+                                d.metadata.update({
+                                    "source": filename,
+                                    "page": i+1,
+                                    "file_type": "pdf",
+                                    "doc_id": f"{filename}_page_{i+1}"
+                                })
+                            documents.extend(docs)
                         
-            except Exception as e:
-                logging.error(f"Lỗi đọc file {filename}: {e}")
-                
-    logging.info(f"Tổng cộng đã load: {len(documents)} tài liệu gốc.")
+                        elif filename.lower().endswith(".docx"):
+                            text = docx2txt.process(file_path)
+                            if text.strip():
+                                doc = Document(
+                                    page_content=text,
+                                    metadata={
+                                        "source": filename,
+                                        "file_type": "docx",
+                                        "doc_id": filename
+                                    }
+                                )
+                                documents.append(doc)
+                        
+                        elif filename.lower().endswith((".xlsx", ".xls", ".csv")):
+                            excel_docs = process_excel_file(file_path, filename)
+                            documents.extend(excel_docs)
+                            
+                        elif filename.lower().endswith((".txt", ".md")):
+                            with open(file_path, "r", encoding="utf-8") as f:
+                                text = f.read()
+                            if text.strip():
+                                doc = Document(
+                                    page_content=text,
+                                    metadata={
+                                        "source": filename,
+                                        "file_type": "txt",
+                                        "doc_id": filename
+                                    }
+                                )
+                                documents.append(doc)
+                                
+                    except Exception as e:
+                        logging.error(f"Lỗi đọc file {filename}: {e}")
+    
+    logging.info(f"Đã load {len(documents)} tài liệu")
     return documents
 
 def get_retriever_chain():
+    """Tạo retriever tối ưu cho Hugging Face"""
     logging.info("--- Tải Embedding Model ---")
-    embedding_model = HuggingFaceEmbeddings(model_name="sentence-transformers/paraphrase-multilingual-MiniLM-L12-v2")
+    
+    # Model tối ưu cho tiếng Việt và memory
+    embedding_model = HuggingFaceEmbeddings(
+        model_name="sentence-transformers/paraphrase-multilingual-MiniLM-L6-v2",
+        model_kwargs={'device': 'cpu'},
+        encode_kwargs={'normalize_embeddings': True}
+    )
     
     vectorstore = None
     splits = []
-
-    chroma_settings = Settings(anonymized_telemetry=False)
-
+    
+    chroma_settings = Settings(
+        anonymized_telemetry=False,
+        allow_reset=True
+    )
+    
     if FORCE_REBUILD_DB and os.path.exists(DB_PATH):
-        logging.warning("Đang xóa DB cũ theo yêu cầu FORCE_REBUILD...")
+        logging.warning("Đang xóa DB cũ...")
         rmtree(DB_PATH, ignore_errors=True)
-
+    
     if os.path.exists(DB_PATH) and os.listdir(DB_PATH):
         try:
             vectorstore = Chroma(
-                persist_directory=DB_PATH, 
+                persist_directory=DB_PATH,
                 embedding_function=embedding_model,
-                client_settings=chroma_settings 
+                client_settings=chroma_settings
             )
             
-            existing_data = vectorstore.get()
-            if existing_data['documents']:
-                for text, meta in zip(existing_data['documents'], existing_data['metadatas']):
+            # Kiểm tra số lượng documents
+            count = vectorstore._collection.count()
+            if count > 0:
+                logging.info(f"Đã khôi phục {count} documents từ DB")
+                # Lấy splits cho BM25
+                results = vectorstore._collection.get()
+                for text, meta in zip(results['documents'], results['metadatas']):
                     splits.append(Document(page_content=text, metadata=meta))
-                logging.info(f"Đã khôi phục {len(splits)} chunks từ DB.")
             else:
-                logging.warning("DB rỗng, sẽ tạo mới.")
+                logging.warning("DB rỗng, tạo mới...")
                 vectorstore = None
         except Exception as e:
-            logging.error(f"DB lỗi: {e}. Đang reset...")
+            logging.error(f"DB lỗi: {e}")
             rmtree(DB_PATH, ignore_errors=True)
             vectorstore = None
-
+    
     if not vectorstore:
-        logging.info("--- Tạo Index dữ liệu mới ---")
+        logging.info("--- Tạo Index mới ---")
         raw_docs = load_documents_from_folder(DATA_PATH)
         if not raw_docs:
-            logging.warning("Không có dữ liệu trong thư mục medical_data.")
+            logging.warning("Không có dữ liệu")
             return None
-
-        text_splitter = RecursiveCharacterTextSplitter(chunk_size=1000, chunk_overlap=200)
+        
+        # Text splitter tối ưu cho y tế
+        text_splitter = RecursiveCharacterTextSplitter(
+            chunk_size=800,  # Tăng độ chính xác với chunk nhỏ hơn
+            chunk_overlap=150,
+            separators=["\n\n", "\n", "。", ".", "!", "?", "；", ";", "，", ",", " ", ""],
+            length_function=len,
+        )
+        
         splits = text_splitter.split_documents(raw_docs)
+        logging.info(f"Đã chia thành {len(splits)} chunks")
+        
+        # Giảm memory bằng cách xóa raw_docs
+        del raw_docs
+        gc.collect()
         
         vectorstore = Chroma.from_documents(
-            documents=splits, 
-            embedding=embedding_model, 
+            documents=splits,
+            embedding=embedding_model,
             persist_directory=DB_PATH,
             client_settings=chroma_settings
         )
-        logging.info("Đã lưu VectorStore thành công.")
-
-    vector_retriever = vectorstore.as_retriever(search_kwargs={"k": 10})
     
+    # Tăng số lượng retrieved documents để bù đắp độ chính xác
+    vector_retriever = vectorstore.as_retriever(
+        search_kwargs={
+            "k": 15,  # Tăng từ 10 lên 15
+            "score_threshold": 0.3  # Ngưỡng similarity
+        }
+    )
+    
+    # BM25 Retriever cho keyword matching
     if splits:
         bm25_retriever = BM25Retriever.from_documents(splits)
-        bm25_retriever.k = 10
-        ensemble_retriever = EnsembleRetriever(
-            retrievers=[bm25_retriever, vector_retriever],
-            weights=[0.4, 0.6] 
-        )
-    else:
-        ensemble_retriever = vector_retriever
-
-    logging.info("--- Tải Reranker Model (BGE-M3) ---")
-    reranker_model = HuggingFaceCrossEncoder(model_name="BAAI/bge-reranker-v2-m3")
-    compressor = CrossEncoderReranker(model=reranker_model, top_n=5)
+        bm25_retriever.k = 15  # Tăng số documents
     
-    final_retriever = ContextualCompressionRetriever(
-        base_compressor=compressor, 
-        base_retriever=ensemble_retriever
+    # Ensemble Retriever với weights tối ưu
+    ensemble_retriever = EnsembleRetriever(
+        retrievers=[bm25_retriever, vector_retriever],
+        weights=[0.5, 0.5]  # Cân bằng giữa keyword và semantic
     )
     
-    return final_retriever
+    # Memory management
+    gc.collect()
+    
+    return ensemble_retriever
 
 class DeepMedBot:
     def __init__(self):
         self.rag_chain = None
         self.ready = False
+        self.retriever = None
+        self.llm = None
+        self.chat_history = []  # Lưu history riêng
         
+        logging.info("Initializing DeepMedBot...")
+    
+    def initialize(self):
+        """Khởi tạo lazy để giảm startup time"""
+        if self.ready:
+            return True
+            
         if not GOOGLE_API_KEY:
-            logging.error("⚠️ Thiếu GOOGLE_API_KEY! Vui lòng thiết lập biến môi trường.")
-            return
-
+            logging.error("⚠️ Thiếu GOOGLE_API_KEY!")
+            return False
+        
         try:
+            # Khởi tạo retriever
             self.retriever = get_retriever_chain()
-            if not self.retriever:
-                logging.warning("⚠️ Chưa có dữ liệu để Retreive. Bot sẽ chỉ trả lời bằng kiến thức nền.")
             
+            # Khởi tạo LLM với config tối ưu
             self.llm = ChatGoogleGenerativeAI(
-                model="gemini-2.5-flash", 
-                temperature=0.11,
-                google_api_key=GOOGLE_API_KEY
+                model="gemini-1.5-flash",  # Dùng flash thay vì 2.5 cho ổn định
+                temperature=0.1,
+                google_api_key=GOOGLE_API_KEY,
+                max_output_tokens=2000,
+                timeout=30
             )
+            
             self._build_chains()
             self.ready = True
-            logging.info("✅ Bot DeepMed đã sẵn sàng phục vụ!")
+            logging.info("✅ DeepMedBot đã sẵn sàng!")
+            return True
+            
         except Exception as e:
-            logging.error(f"🔥 Lỗi khởi tạo bot: {e}")
-            logging.debug(traceback.format_exc())
-
+            logging.error(f"🔥 Lỗi khởi tạo: {e}")
+            return False
+    
     def _build_chains(self):
-        context_system_prompt = (
-            "Dựa trên lịch sử chat và câu hỏi mới nhất của người dùng, "
-            "hãy viết lại câu hỏi đó thành một câu đầy đủ ngữ cảnh để hệ thống có thể hiểu được. "
-            "KHÔNG trả lời câu hỏi, chỉ viết lại nó."
+        """Xây dựng chains với prompt tối ưu"""
+        # Contextualize question với medical focus
+        contextualize_q_system_prompt = (
+            "Bạn là trợ lý y tế. Dựa vào lịch sử chat và câu hỏi mới, "
+            "hãy viết lại câu hỏi thành một phiên bản đầy đủ, rõ ràng, "
+            "chuyên nghiệp về y tế để tìm kiếm thông tin.\n"
+            "Ví dụ:\n"
+            "User: 'tác dụng phụ?' -> 'Thuốc này có những tác dụng phụ gì?'\n"
+            "KHÔNG trả lời câu hỏi, chỉ VIẾT LẠI câu hỏi."
         )
-        context_prompt = ChatPromptTemplate.from_messages([
-            ("system", context_system_prompt),
+        
+        contextualize_q_prompt = ChatPromptTemplate.from_messages([
+            ("system", contextualize_q_system_prompt),
             MessagesPlaceholder("chat_history"),
             ("human", "{input}"),
         ])
         
-        if self.retriever:
-            history_aware_retriever = create_history_aware_retriever(
-                self.llm, self.retriever, context_prompt
-            )
-        
-        qa_system_prompt = (
-            "Bạn là 'DeepMed-AI' - Trợ lý Dược lâm sàng tại Trung Tâm Y Tế. "
-            "Sử dụng các thông tin được cung cấp trong phần Context dưới đây để trả lời câu hỏi về thuốc, bệnh học và y lệnh.\n"
-            "Nếu Context có dữ liệu từ Excel, hãy trình bày dạng bảng hoặc gạch đầu dòng rõ ràng.\n"
-            "Nếu không tìm thấy thông tin trong Context, hãy nói 'Tôi không tìm thấy thông tin trong dữ liệu nội bộ' và gợi ý dựa trên kiến thức y khoa chung của bạn.\n\n"
-            "Context:\n{context}"
-        )
+        # QA prompt tối ưu cho y tế
+        qa_system_prompt = """
+Bạn là "DeepMed AI" - Trợ lý Dược lâm sàng thông minh tại Trung Tâm Y Tế Khu Vực Thanh Ba.
+
+HƯỚNG DẪN TRẢ LỜI:
+1. **NGUYÊN TẮC VÀNG**: Luôn kiểm tra kỹ thông tin từ Context trước khi trả lời
+2. **ĐỊNH DẠNG RÕ RÀNG**:
+   - Thuốc: Tên thuốc (IN HOA), liều lượng, chống chỉ định, tác dụng phụ
+   - Bệnh nhân: Mã BN, tuổi, chẩn đoán, phác đồ
+   - Số liệu: Trình bày dạng bảng hoặc bullet points
+3. **MỨC ĐỘ TIN CẬY**:
+   ✅ "Theo dữ liệu nội bộ: [thông tin]"
+   ⚠️  "Thông tin không đầy đủ trong dữ liệu, theo kiến thức y khoa: [thông tin]"
+   ❌ "Không tìm thấy trong dữ liệu, vui lòng kiểm tra lại"
+4. **AN TOÀN Y TẾ**: Luôn nhắc "Vui lòng tham khảo ý kiến bác sĩ trước khi sử dụng"
+
+Context:
+{context}
+
+Hãy trả lời câu hỏi dựa trên Context trên. Nếu không có thông tin trong Context, hãy:
+1. Nói rõ "Không tìm thấy trong dữ liệu nội bộ"
+2. Cung cấp kiến thức y khoa chung (nếu có)
+3. Gợi ý tham khảo bác sĩ chuyên khoa
+
+Câu hỏi: {input}
+"""
         
         qa_prompt = ChatPromptTemplate.from_messages([
             ("system", qa_system_prompt),
@@ -238,93 +347,304 @@ class DeepMedBot:
             ("human", "{input}"),
         ])
         
-        question_answer_chain = create_stuff_documents_chain(self.llm, qa_prompt)
-        
         if self.retriever:
+            # Tạo history-aware retriever
+            history_aware_retriever = create_history_aware_retriever(
+                self.llm, self.retriever, contextualize_q_prompt
+            )
+            
+            # Tạo chain trả lời
+            question_answer_chain = create_stuff_documents_chain(self.llm, qa_prompt)
+            
+            # Tạo retrieval chain hoàn chỉnh
             self.rag_chain = create_retrieval_chain(history_aware_retriever, question_answer_chain)
         else:
+            # Fallback chain
             self.rag_chain = qa_prompt | self.llm
-
+    
     def chat_stream(self, message: str, history: list):
+        """Stream chat với memory management"""
         if not self.ready:
-            yield "Hệ thống đang khởi động hoặc gặp lỗi cấu hình."
-            return
-
+            if not self.initialize():
+                yield "Hệ thống đang khởi động..."
+                return
+        
+        # Giới hạn history để tránh memory leak
         chat_history = []
         for u, b in history[-MAX_HISTORY_TURNS:]:
             chat_history.append(HumanMessage(content=str(u)))
             chat_history.append(AIMessage(content=str(b)))
-
+        
         full_response = ""
         retrieved_docs = []
         
         try:
-            stream_input = {"input": message, "chat_history": chat_history} if self.retriever else {"input": message, "chat_history": chat_history}
+            # Thêm delay nhỏ để tránh timeout
+            import time
             
-            if self.rag_chain:
-                for chunk in self.rag_chain.stream(stream_input):
-                    
+            # Stream response
+            if hasattr(self.rag_chain, 'stream'):
+                for chunk in self.rag_chain.stream({
+                    "input": message,
+                    "chat_history": chat_history
+                }):
                     if isinstance(chunk, dict):
                         if "answer" in chunk:
                             full_response += chunk["answer"]
-                            yield full_response 
-                        
+                            yield full_response
                         if "context" in chunk:
                             retrieved_docs = chunk["context"]
-
                     elif hasattr(chunk, 'content'):
                         full_response += chunk.content
                         yield full_response
-                    
-                    elif isinstance(chunk, str):
-                        full_response += chunk
-                        yield full_response
-
-                if retrieved_docs:
-                    refs = self._build_references_text(retrieved_docs)
-                    if refs:
-                        full_response += f"\n\n---\n📚 **Nguồn tham khảo:**\n{refs}"
-                        yield full_response
-                        
+                    time.sleep(0.01)  # Small delay
+            else:
+                # Fallback non-stream
+                response = self.rag_chain.invoke({
+                    "input": message,
+                    "chat_history": chat_history
+                })
+                full_response = response.content if hasattr(response, 'content') else str(response)
+                yield full_response
+            
+            # Thêm references nếu có
+            if retrieved_docs:
+                refs = self._build_references_text(retrieved_docs)
+                if refs:
+                    full_response += f"\n\n---\n📚 **Tài liệu tham khảo:**\n{refs}"
+                    yield full_response
+            
+            # Memory cleanup
+            gc.collect()
+            
         except Exception as e:
-            logging.error(f"Lỗi khi chat: {e}")
-            logging.debug(traceback.format_exc())
-            yield f"Đã xảy ra lỗi: {str(e)}"
-
+            logging.error(f"Chat error: {e}")
+            yield f"⚠️ Có lỗi xảy ra: {str(e)[:100]}"
+    
     @staticmethod
     def _build_references_text(docs) -> str:
-        lines = []
-        seen = set()
+        """Xây dựng references với format đẹp"""
+        references = {}
         for doc in docs:
-            src = doc.metadata.get("source", "Tài liệu")
-            row_info = ""
-            if "row" in doc.metadata:
-                row_info = f"(Dòng {doc.metadata['row']})"
+            source = doc.metadata.get("source", "Tài liệu")
+            file_type = doc.metadata.get("file_type", "")
+            row = doc.metadata.get("row", "")
             
-            ref_str = f"- {src} {row_info}"
-            
-            if ref_str not in seen:
-                lines.append(ref_str)
-                seen.add(ref_str)
-        return "\n".join(lines)
+            key = f"{source}_{row}"
+            if key not in references:
+                ref_info = f"📄 {source}"
+                if file_type:
+                    ref_info += f" ({file_type.upper()})"
+                if row:
+                    ref_info += f" - Dòng {row}"
+                references[key] = ref_info
+        
+        return "\n".join(references.values())
 
+# Global bot instance với lazy loading
 bot = DeepMedBot()
 
 def gradio_chat_stream(message, history):
+    """Wrapper cho Gradio"""
     yield from bot.chat_stream(message, history)
 
+# CSS responsive cho cả mobile và PC
 css = """
-.gradio-container {min_height: 600px !important;}
-h1 {text-align: center; color: #2E86C1;}
-"""
+/* Base styles */
+.gradio-container {
+    min-height: 100vh !important;
+    max-width: 100% !important;
+    margin: 0 auto !important;
+    padding: 10px !important;
+}
+
+/* Header */
+h1 {
+    text-align: center;
+    color: #2E86C1;
+    font-size: 24px !important;
+    margin: 10px 0 !important;
+    padding: 10px !important;
+}
+
+/* Chat container */
+#chatbot {
+    min-height: 400px !important;
+    max-height: 60vh !important;
+    overflow-y: auto !important;
+    border: 1px solid #e0e0e0 !important;
+    border-radius: 10px !important;
+    padding: 15px !important;
+    background: #f9f9f9 !important;
+}
 
-with gr.Blocks(css=css, title="DeepMed AI") as demo:
-    gr.Markdown("# 🏥 DeepMed AI - Trợ lý Lâm Sàng")
-    gr.Markdown("Hệ thống hỗ trợ lâm sàng tại Trung Tâm Y Tế Khu Vực Thanh Ba.")
+/* Messages */
+.user, .assistant {
+    padding: 10px 15px !important;
+    margin: 8px 0 !important;
+    border-radius: 15px !important;
+    max-width: 85% !important;
+    word-wrap: break-word !important;
+}
+
+.user {
+    background: #E3F2FD !important;
+    margin-left: auto !important;
+}
+
+.assistant {
+    background: #F5F5F5 !important;
+    margin-right: auto !important;
+}
+
+/* Input area */
+#text-input {
+    border-radius: 20px !important;
+    padding: 12px 20px !important;
+    font-size: 14px !important;
+    border: 2px solid #2E86C1 !important;
+}
+
+/* Buttons */
+button {
+    border-radius: 20px !important;
+    padding: 10px 20px !important;
+    font-weight: bold !important;
+    transition: all 0.3s !important;
+}
+
+button:hover {
+    transform: translateY(-2px) !important;
+    box-shadow: 0 4px 8px rgba(0,0,0,0.1) !important;
+}
+
+/* Footer */
+.footer {
+    text-align: center;
+    padding: 10px;
+    color: #666;
+    font-size: 12px;
+}
+
+/* Mobile optimization */
+@media screen and (max-width: 768px) {
+    .gradio-container {
+        padding: 5px !important;
+    }
     
-    chat_interface = gr.ChatInterface(
-        fn=gradio_chat_stream,
-    )
+    h1 {
+        font-size: 20px !important;
+        padding: 5px !important;
+    }
+    
+    #chatbot {
+        min-height: 300px !important;
+        max-height: 50vh !important;
+        padding: 10px !important;
+    }
+    
+    .user, .assistant {
+        max-width: 90% !important;
+        padding: 8px 12px !important;
+        font-size: 14px !important;
+    }
+    
+    #text-input {
+        padding: 10px 15px !important;
+        font-size: 13px !important;
+    }
+    
+    button {
+        padding: 8px 16px !important;
+        margin: 5px !important;
+    }
+}
+
+/* PC optimization */
+@media screen and (min-width: 1200px) {
+    .gradio-container {
+        max-width: 900px !important;
+    }
+    
+    #chatbot {
+        max-height: 500px !important;
+    }
+}
+
+/* Loading animation */
+@keyframes pulse {
+    0% { opacity: 1; }
+    50% { opacity: 0.5; }
+    100% { opacity: 1; }
+}
+
+.typing {
+    animation: pulse 1.5s infinite;
+}
+
+/* Table formatting for medical data */
+table {
+    border-collapse: collapse;
+    width: 100%;
+    margin: 10px 0;
+}
+
+th, td {
+    border: 1px solid #ddd;
+    padding: 8px;
+    text-align: left;
+}
+
+th {
+    background-color: #f2f2f2;
+}
+
+/* Scrollbar styling */
+::-webkit-scrollbar {
+    width: 6px;
+}
+
+::-webkit-scrollbar-track {
+    background: #f1f1f1;
+}
+
+::-webkit-scrollbar-thumb {
+    background: #888;
+    border-radius: 3px;
+}
+
+::-webkit-scrollbar-thumb:hover {
+    background: #555;
+}
+"""
+
+# Config cho Hugging Face Spaces
+def get_spaces_config():
+    return {
+        "title": "DeepMed AI - Medical Assistant",
+        "description": "Trợ lý lâm sàng AI cho trung tâm y tế",
+        "thumbnail": "https://huggingface.co/spaces/your-space/your-app/raw/main/thumbnail.png",
+        "theme": "light",
+        "sdk": "gradio",
+        "sdk_version": "4.0.0",
+    }
+
+# Memory management
+def cleanup():
+    """Cleanup function for Hugging Face"""
+    if torch.cuda.is_available():
+        torch.cuda.empty_cache()
+    gc.collect()
 
 if __name__ == "__main__":
-    demo.launch()
+    # Hugging Face Spaces config
+    demo.queue(max_size=20)  # Giới hạn queue
+    demo.launch(
+        server_name="0.0.0.0",
+        server_port=7860,
+        show_error=True,
+        debug=False,
+        share=False,
+        favicon_path=None
+    )
+    cleanup()