Upload app.py

app.py

@@ -12,7 +12,7 @@ import json
 from huggingface_hub import hf_hub_download
 
 # ==========================================
-# 1. ĐỊNH NGHĨA KIẾN TRÚC MODEL (GraphSAGE)
+# 1. ĐỊNH NGHĨA MODEL
 # ==========================================
 class SAGE(nn.Module):
     def __init__(self, in_dim, h=128, out_dim=2, p_drop=0.3):
@@ -25,251 +25,271 @@ class SAGE(nn.Module):
         self.drop = nn.Dropout(p_drop)
 
     def forward(self, x, edge_index):
-        # Layer 1
         x = self.conv1(x, edge_index)
         x = self.bn1(x)
         x = F.relu(x)
         x = self.drop(x)
-        # Layer 2
         x = self.conv2(x, edge_index)
         x = self.bn2(x)
         x = F.relu(x)
         x = self.drop(x)
-        # Output
         return self.head(x)
 
 # ==========================================
-# 2. QUẢN LÝ DỮ LIỆU & TẢI MODEL
+# 2. QUẢN LÝ RESOURCE (DATA & MODEL)
 # ==========================================
 REPO_ID = "uyen1109/eth-fraud-gnn-uyenuyen-v3"
 TOKEN = os.getenv("HF_TOKEN")
 
-# Biến toàn cục lưu trữ dữ liệu
+# Danh sách feature mặc định (Fallback) để UI không bị trống nếu lỗi load file
+DEFAULT_FEATURES = [
+    'out_deg', 'in_deg', 'eth_out_sum', 'eth_in_sum', 
+    'unique_dst_cnt', 'unique_src_cnt', 'first_seen_ts', 'last_seen_ts',
+    'pr', 'clust_coef', 'betw', 'feat_11', 'feat_12', 'feat_13', 'feat_14'
+]
+
 GLOBAL_DATA = {
     "model": None,
     "df_scores": pd.DataFrame(),
     "df_edges": pd.DataFrame(),
-    "feature_cols": [],
-    "input_dim": 0
+    "feature_cols": DEFAULT_FEATURES, # Luôn có giá trị mặc định
+    "status": "Initializing..."
 }
 
 def smart_load_file(filename):
-    """Tìm file ở root hoặc thư mục hf_export"""
-    possible_paths = [f"hf_export/{filename}", filename]
-    for p in possible_paths:
+    """Thử tải file từ repo, ưu tiên hf_export"""
+    paths = [f"hf_export/{filename}", filename]
+    for p in paths:
         try:
+            # Thử tải với token trước, nếu lỗi thử không token (public repo)
             return hf_hub_download(repo_id=REPO_ID, filename=p, token=TOKEN)
         except:
-            continue
+            try:
+                return hf_hub_download(repo_id=REPO_ID, filename=p, token=None)
+            except:
+                continue
     return None
 
 def load_resources():
-    print("⏳ Loading resources...")
+    logs = []
+    print("⏳ Starting Resource Loading...")
     
-    # 1. Load Scores (Lookup Data)
+    # 1. Load Scores
     try:
         path = smart_load_file("scores/node_scores_with_labels.csv") or smart_load_file("node_scores_with_labels.csv")
         if path:
             df = pd.read_csv(path)
-            if "address" in df.columns:
-                df["address"] = df["address"].astype(str).str.lower().str.strip()
-                df.set_index("address", inplace=True)
+            # Chuẩn hóa cột địa chỉ: tìm cột chứa chữ 'address' hoặc 'id'
+            addr_col = next((c for c in df.columns if 'addr' in c.lower() or 'id' in c.lower()), df.columns[0])
+            df[addr_col] = df[addr_col].astype(str).str.lower().str.strip()
+            # Đặt index là địa chỉ để tra cứu nhanh
+            df.set_index(addr_col, inplace=True)
             GLOBAL_DATA["df_scores"] = df
-            print(f"✅ Loaded {len(df)} scores.")
+            msg = f"✅ Loaded Scores: {len(df)} rows (Index col: {addr_col})"
+            print(msg)
+            logs.append(msg)
+        else:
+            logs.append("⚠️ Scores CSV not found.")
     except Exception as e:
-        print(f"⚠️ Score load error: {e}")
+        logs.append(f"❌ Error loading scores: {str(e)}")
 
-    # 2. Load Edges (Graph Visualization)
+    # 2. Load Edges
     try:
         path = smart_load_file("graph/edges_all.csv") or smart_load_file("edges_all.csv")
         if path:
-            # Chỉ load cột cần thiết để tiết kiệm RAM
-            GLOBAL_DATA["df_edges"] = pd.read_csv(path, usecols=["src", "dst", "edge_type"])
-            GLOBAL_DATA["df_edges"]["src"] = GLOBAL_DATA["df_edges"]["src"].astype(str).str.lower().str.strip()
-            GLOBAL_DATA["df_edges"]["dst"] = GLOBAL_DATA["df_edges"]["dst"].astype(str).str.lower().str.strip()
-            print(f"✅ Loaded edges.")
-    except Exception as e:
-        print(f"⚠️ Edge load error: {e}")
+            df = pd.read_csv(path, usecols=["src", "dst", "edge_type"])
+            df["src"] = df["src"].astype(str).str.lower().str.strip()
+            df["dst"] = df["dst"].astype(str).str.lower().str.strip()
+            GLOBAL_DATA["df_edges"] = df
+            print("✅ Loaded Edges.")
+    except:
+        print("⚠️ Edges CSV not found (Graph viz will be disabled).")
 
-    # 3. Load Model Weights & Detect Dimension
+    # 3. Load Model
     try:
         model_path = smart_load_file("pytorch_model.bin")
         if model_path:
             state_dict = torch.load(model_path, map_location=torch.device('cpu'))
-            # Tự động phát hiện input dimension từ trọng số
+            # Tự động phát hiện input dim
             detected_dim = state_dict['conv1.lin_l.weight'].shape[1]
-            GLOBAL_DATA["input_dim"] = detected_dim
             
-            # Khởi tạo model
             model = SAGE(in_dim=detected_dim, h=128, out_dim=2, p_drop=0.3)
             model.load_state_dict(state_dict)
             model.eval()
             GLOBAL_DATA["model"] = model
-            print(f"✅ Model loaded (Input Dim: {detected_dim})")
             
-            # Load tên cột features (nếu có)
+            # Cập nhật danh sách feature cột nếu có
             cols_path = smart_load_file("feature_columns.json")
             if cols_path:
                 with open(cols_path, 'r') as f:
                     cols = json.load(f)
-                    # Cắt hoặc thêm cho khớp với detected_dim
-                    if len(cols) > detected_dim:
+                    # Điều chỉnh cho khớp detected_dim
+                    if len(cols) >= detected_dim:
                         GLOBAL_DATA["feature_cols"] = cols[:detected_dim]
-                    elif len(cols) < detected_dim:
-                        GLOBAL_DATA["feature_cols"] = cols + [f"Feat_{i}" for i in range(detected_dim - len(cols))]
                     else:
-                        GLOBAL_DATA["feature_cols"] = cols
+                        GLOBAL_DATA["feature_cols"] = cols + [f"F_{i}" for i in range(len(cols), detected_dim)]
             else:
+                # Nếu không có file json, tạo dummy name cho đủ số lượng
                 GLOBAL_DATA["feature_cols"] = [f"Feature_{i}" for i in range(detected_dim)]
+                
+            logs.append(f"✅ Model Loaded (Input Dim: {detected_dim})")
+        else:
+            logs.append("❌ pytorch_model.bin not found.")
     except Exception as e:
-        print(f"❌ Critical Model Error: {e}")
+        logs.append(f"❌ Model Load Error: {str(e)}")
+
+    GLOBAL_DATA["status"] = "\n".join(logs)
 
-# Gọi hàm load khi khởi động
+# Chạy load ngay lập tức
 load_resources()
 
 # ==========================================
-# 3. LOGIC XỬ LÝ (HANDLERS)
+# 3. LOGIC XỬ LÝ
 # ==========================================
 
-# --- Xử lý Tab 1: Tra cứu ---
-def draw_ego_graph(address):
-    """Vẽ đồ thị mạng lưới giao dịch cục bộ"""
+def draw_graph(address):
     df = GLOBAL_DATA["df_edges"]
     if df.empty: return None
     
-    # Lấy 20 giao dịch gần nhất
-    subset = df[(df["src"] == address) | (df["dst"] == address)].head(20)
+    # Tìm giao dịch liên quan (cả in và out)
+    subset = df[(df["src"] == address) | (df["dst"] == address)].head(30)
     if subset.empty: return None
 
     G = nx.from_pandas_edgelist(subset, "src", "dst", edge_attr="edge_type", create_using=nx.DiGraph())
-    pos = nx.spring_layout(G, seed=42, k=0.9)
     
     plt.figure(figsize=(8, 8))
-    colors = ["#FF4500" if n == address else "#1E90FF" for n in G.nodes()]
-    sizes = [400 if n == address else 100 for n in G.nodes()]
+    pos = nx.spring_layout(G, k=0.8, seed=42)
     
-    nx.draw_networkx_nodes(G, pos, node_color=colors, node_size=sizes, alpha=0.9)
-    nx.draw_networkx_edges(G, pos, alpha=0.4, arrowstyle='->')
-    nx.draw_networkx_labels(G, pos, labels={n: n[:4] for n in G.nodes()}, font_size=8)
+    # Màu sắc: Target màu đỏ, Neighbor màu xanh
+    node_colors = ["#FF4500" if n == address else "#1E90FF" for n in G.nodes()]
+    node_sizes = [400 if n == address else 150 for n in G.nodes()]
     
-    plt.title(f"Transaction Graph: {address[:6]}...")
+    nx.draw_networkx_nodes(G, pos, node_color=node_colors, node_size=node_sizes, alpha=0.9)
+    nx.draw_networkx_edges(G, pos, alpha=0.3, arrowstyle='->')
+    # Label ngắn gọn
+    nx.draw_networkx_labels(G, pos, labels={n: n[:4] + ".." for n in G.nodes()}, font_size=8)
+    
+    plt.title(f"Ego Graph: {address[:6]}...")
     plt.axis('off')
     return plt.gcf()
 
-def lookup_address(address):
-    address = address.lower().strip()
+def lookup_handler(address):
+    # Chuẩn hóa input cực mạnh để khớp với index CSV
+    raw_addr = str(address).strip().lower()
+    
     df = GLOBAL_DATA["df_scores"]
+    found_row = None
     
-    # 1. Tìm trong Database
-    if not df.empty and address in df.index:
-        row = df.loc[address]
-        score = float(row.get("prob_criminal", row.get("susp", 0.0)))
-        label_val = row.get("label", -1)
-        
-        status = "CRIMINAL 🔴" if score > 0.5 else "BENIGN 🟢"
-        if label_val == 1: status += " (Verified Criminal)"
-        elif label_val == 0: status += " (Verified Benign)"
+    # Thử các trường hợp khớp
+    if raw_addr in df.index:
+        found_row = df.loc[raw_addr]
+    elif raw_addr.startswith("0x") and raw_addr[2:] in df.index: # Thử bỏ 0x
+        found_row = df.loc[raw_addr[2:]]
+    elif ("0x" + raw_addr) in df.index: # Thử thêm 0x
+        found_row = df.loc["0x" + raw_addr]
         
-        plot = draw_ego_graph(address)
-        return (
-            f"### ✅ Found in Database\n**Score:** {score:.4f}\n**Status:** {status}",
-            plot,
-            gr.update(visible=False) # Ẩn thông báo chuyển tab
-        )
-    
-    # 2. Nếu không tìm thấy
-    msg = (
-        f"### ❌ Address Not Found\n"
-        f"Địa chỉ `{address}` chưa được model chấm điểm trước đó.\n"
-        f"Vui lòng chuyển sang tab **'Inductive Prediction'** để nhập thông số và dự đoán."
+    if found_row is not None:
+        # Lấy điểm số
+        try:
+            # Xử lý trường hợp duplicate index hoặc series
+            if isinstance(found_row, pd.DataFrame): 
+                found_row = found_row.iloc[0]
+            
+            score = float(found_row.get("prob_criminal", found_row.get("susp", 0.0)))
+            label = int(found_row.get("label", -1))
+            
+            risk = "CRITICAL 🔴" if score > 0.8 else ("HIGH 🟠" if score > 0.5 else "LOW 🟢")
+            label_text = "Unknown"
+            if label == 1: label_text = "Criminal (True Label)"
+            elif label == 0: label_text = "Benign (True Label)"
+            
+            info = (
+                f"### ✅ Address Found\n"
+                f"- **Risk Score:** {score:.4f}\n"
+                f"- **Risk Level:** {risk}\n"
+                f"- **Dataset Label:** {label_text}"
+            )
+            return info, draw_graph(raw_addr)
+        except Exception as e:
+            return f"Error parsing row: {e}", None
+
+    # Nếu không tìm thấy
+    return (
+        f"### ❌ Not Found in Database\n"
+        f"Address `{raw_addr}` does not exist in `node_scores_with_labels.csv`.\n"
+        f"Please verify the address or use the **Inductive Prediction** tab.", 
+        None
     )
-    return msg, None, gr.update(visible=True)
 
-# --- Xử lý Tab 2: Dự đoán ---
-def predict_manual(*features):
+def predict_handler(*features):
     model = GLOBAL_DATA["model"]
     if model is None:
-        return "❌ Model not loaded."
+        return f"❌ Model failed to load properly.\n\nLogs:\n{GLOBAL_DATA['status']}"
     
     try:
-        # Chuyển input thành tensor
-        feat_vals = [float(f) for f in features]
-        x = torch.tensor([feat_vals], dtype=torch.float)
-        # Edge index giả (node cô lập)
+        x = torch.tensor([[float(f) for f in features]], dtype=torch.float)
         edge_index = torch.tensor([[], []], dtype=torch.long)
         
         with torch.no_grad():
             logits = model(x, edge_index)
-            probs = torch.softmax(logits, dim=1)
-            prob_crim = probs[0][1].item()
+            prob = torch.softmax(logits, dim=1)[0][1].item()
             
-        label = "CRIMINAL 🔴" if prob_crim > 0.5 else "BENIGN 🟢"
+        verdict = "CRIMINAL 🔴" if prob > 0.5 else "BENIGN 🟢"
         return (
-            f"### 🧠 Inductive Prediction Result\n"
-            f"- **Fraud Probability:** {prob_crim*100:.2f}%\n"
-            f"- **Model Verdict:** {label}\n"
-            f"*(Dự đoán dựa trên {len(feat_vals)} features đầu vào)*"
+            f"### 🧠 Prediction Result\n"
+            f"- **Fraud Probability:** {prob*100:.2f}%\n"
+            f"- **Verdict:** {verdict}"
         )
     except Exception as e:
-        return f"⚠️ Error: {str(e)}"
+        return f"Prediction Error: {str(e)}"
 
 # ==========================================
-# 4. GIAO DIỆN GRADIO (UI)
+# 4. UI SETUP
 # ==========================================
-with gr.Blocks(title="ETH Fraud GNN Hybrid") as demo:
-    gr.Markdown("# 🕵️‍♀️ Ethereum Fraud Detection System (Hybrid)")
-    gr.Markdown("Kết hợp tra cứu dữ liệu lịch sử và khả năng dự đoán (Inductive) trên dữ liệu mới.")
+with gr.Blocks(title="ETH Fraud GNN") as demo:
+    gr.Markdown("# 🕵️‍♀️ Ethereum Fraud GNN (Hybrid V3)")
     
+    # Hiển thị trạng thái load hệ thống (ẩn đi nếu muốn gọn)
+    with gr.Accordion("System Status / Logs", open=False):
+        gr.Markdown(GLOBAL_DATA["status"])
+
     with gr.Tabs():
-        # === TAB 1: TRA CỨU ===
+        # TAB 1: LOOKUP
         with gr.TabItem("🔍 Lookup Address"):
             with gr.Row():
-                with gr.Column():
-                    inp_addr = gr.Textbox(label="Ethereum Address", placeholder="0x...")
-                    btn_lookup = gr.Button("Search", variant="primary")
-                with gr.Column():
-                    out_lookup_text = gr.Markdown()
-                    out_plot = gr.Plot(label="Graph")
-                    # Thông báo hướng dẫn chuyển tab (ẩn mặc định)
-                    notice_box = gr.Markdown("👉 **Tip:** Use the 'Inductive Prediction' tab to predict unknown addresses.", visible=False)
-
-            btn_lookup.click(
-                lookup_address, 
-                inputs=inp_addr, 
-                outputs=[out_lookup_text, out_plot, notice_box]
-            )
+                inp_addr = gr.Textbox(label="Enter Address", placeholder="0x...")
+                btn_search = gr.Button("Search", variant="primary")
+            
+            with gr.Row():
+                out_info = gr.Markdown()
+                out_plot = gr.Plot()
+            
+            btn_search.click(lookup_handler, inputs=inp_addr, outputs=[out_info, out_plot])
 
-        # === TAB 2: DỰ ĐOÁN THỦ CÔNG ===
+        # TAB 2: INDUCTIVE
         with gr.TabItem("🧠 Inductive Prediction"):
-            gr.Markdown("### Predict New/Unknown Address")
-            gr.Markdown("Nhập các chỉ số feature của ví để model dự đoán rủi ro.")
+            gr.Markdown("### Predict New Address")
+            gr.Markdown("Enter extracted features manually:")
             
-            input_comps = []
-            if GLOBAL_DATA["model"]:
-                # Tạo lưới nhập liệu động
-                cols = GLOBAL_DATA["feature_cols"]
-                with gr.Row():
-                    # Chia làm 3 cột cho gọn
-                    with gr.Column():
-                        for c in cols[:len(cols)//3]:
-                            input_comps.append(gr.Number(label=c, value=0.0))
-                    with gr.Column():
-                        for c in cols[len(cols)//3 : 2*len(cols)//3]:
-                            input_comps.append(gr.Number(label=c, value=0.0))
-                    with gr.Column():
-                        for c in cols[2*len(cols)//3:]:
-                            input_comps.append(gr.Number(label=c, value=0.0))
-                
-                btn_predict = gr.Button("Run GraphSAGE Inference", variant="primary")
-                out_pred_text = gr.Markdown()
+            # TẠO INPUT ĐỘNG: Dù model có load được hay không, UI vẫn sẽ render dựa trên GLOBAL_DATA["feature_cols"]
+            # Điều này fix lỗi giao diện trống trơn.
+            feat_inputs = []
+            cols = GLOBAL_DATA["feature_cols"]
+            
+            # Chia layout thành 3 cột
+            with gr.Row():
+                col1, col2, col3 = gr.Column(), gr.Column(), gr.Column()
                 
-                btn_predict.click(
-                    predict_manual,
-                    inputs=input_comps,
-                    outputs=out_pred_text
-                )
-            else:
-                gr.Error("Model failed to load. Cannot render inputs.")
+                # Phân phối input vào 3 cột
+                for i, c in enumerate(cols):
+                    target_col = col1 if i % 3 == 0 else (col2 if i % 3 == 1 else col3)
+                    with target_col:
+                        feat_inputs.append(gr.Number(label=c, value=0.0))
+
+            btn_predict = gr.Button("Run Inference", variant="primary")
+            out_pred = gr.Markdown()
+            
+            btn_predict.click(predict_handler, inputs=feat_inputs, outputs=out_pred)
 
 if __name__ == "__main__":
     demo.launch()