Update src/streamlit_app.py

src/streamlit_app.py

@@ -7,27 +7,22 @@ import faiss
 import uuid
 import time
 import sys
-
-# === HuggingFace 模型相關套件 (新增) ===
+# === HuggingFace 模型相關套件 (替換為 InferenceClient) ===
 try:
-    # 確保只在需要時載入，避免在無 GPU 環境下強制載入導致錯誤
-    from transformers import AutoModelForCausalLM, AutoTokenizer, pipeline
-    import torch
-    # 針對本地大模型：
-    # from accelerate import Accelerator # 建議安裝
-    # import bitsandbytes # 建議安裝
+    from huggingface_hub import InferenceClient
+    # 移除本地模型相關導入
+    # from transformers import AutoModelForCausalLM, AutoTokenizer, pipeline
+    # import torch
 except ImportError:
-    st.error("請檢查是否安裝了所有 Hugging Face 相關依賴：pip install transformers torch accelerate bitsandbytes")
-    # 如果缺少，則退出或將相關變數設為 None
-    AutoModelForCausalLM, AutoTokenizer, pipeline, torch = None, None, None, None
-
+    st.error("請檢查是否安裝了所有 Hugging Face 相關依賴：pip install huggingface-hub")
+    # InferenceClient = None # 保留 InferenceClient
+    
 # === LangChain/RAG 相關套件 (保持不變) ===
 from langchain_community.embeddings import HuggingFaceEmbeddings
 from langchain_core.documents import Document
 from langchain_community.vectorstores import FAISS
 from langchain_community.vectorstores.utils import DistanceStrategy
 from langchain_community.docstore.in_memory import InMemoryDocstore
-
 # 嘗試匯入 pypdf
 try:
     import pypdf
@@ -36,11 +31,11 @@ except ImportError:
 
 # --- 頁面設定 ---
 st.set_page_config(page_title="Cybersecurity AI Assistant (Hugging Face RAG & Batch Analysis)", page_icon="🛡️", layout="wide")
-st.title("🛡️ Foundation-Sec-1.1-8B-Instruct with FAISS RAG & Batch Analysis")
-st.markdown("已啟用：**IndexFlatIP** + **L2 正規化** + **Hugging Face LLM**。上傳 JSON 執行批量分析，上傳其他檔案作為 RAG 參考庫。")
+st.title("🛡️ Meta-Llama-3-8B-Instruct with FAISS RAG & Batch Analysis (Inference Client)")
+st.markdown("已啟用：**IndexFlatIP** + **L2 正規化** + **Hugging Face Inference Client (API)**。上傳 JSON 執行批量分析，上傳其他檔案作為 RAG 參考庫。")
 
-# 設定模型 ID (替換為 Hugging Face 模型名稱)
-MODEL_ID = "Qwen/Qwen2.5-7B-Instruct" 
+# 設定模型 ID (替換為您指定的模型)
+MODEL_ID = "fdtn-ai/Foundation-Sec-8B-Instruct" 
 WINDOW_SIZE = 8
 
 # --- 側邊欄設定 ---
@@ -48,8 +43,12 @@ with st.sidebar:
     st.header("⚙️ 設定")
     
     # === 替換為 Hugging Face 模型名稱顯示 (移除 API Key 輸入) ===
-    st.info(f"LLM 模型：**{MODEL_ID}** (Hugging Face Model)")
-    st.warning("⚠️ **注意**: 8B 模型需要大量 RAM/VRAM 和算力。運行可能較慢或失敗。")
+    # ⚠️ 注意: HF Token 必須在環境變數 HF_TOKEN 中設定
+    if not os.environ.get("HF_TOKEN"):
+         st.error("環境變數 **HF_TOKEN** 未設定。請設定後重新啟動應用程式。")
+    
+    st.info(f"LLM 模型：**{MODEL_ID}** (Hugging Face Inference API)")
+    st.warning("⚠️ **注意**: 該模型使用 Inference API 呼叫，請確保您的 HF Token 具有存取權限。")
     
     st.divider()
     
@@ -78,7 +77,10 @@ with st.sidebar:
     
     if json_uploaded_file: # 移除 API Key 檢查
         if st.button("🚀 執行批量分析"):
-            st.session_state.execute_batch_analysis = True
+            if not os.environ.get("HF_TOKEN"):
+                st.error("無法執行，環境變數 **HF_TOKEN** 未設定。")
+            else:
+                 st.session_state.execute_batch_analysis = True
     else:
         st.info("請上傳 JSON 檔案以啟用批量分析按鈕。")
         
@@ -93,7 +95,8 @@ with st.sidebar:
     st.divider()
     
     st.subheader("模型參數")
-    system_prompt = st.text_area("System Prompt (LLM 使用)", value="You are a Senior Security Analyst. Be professional.", height=100)
+    # Llama 3 使用 'system' 角色
+    system_prompt = st.text_area("System Prompt (LLM 使用)", value="You are a Senior Security Analyst, named Ernest. You provide expert, authoritative, and concise advice on Information Security, Network Security, and Cyber Threat Intelligence. Your analysis must be based strictly on the provided context.", height=100)
     max_output_tokens = st.slider("Max Output Tokens", 128, 4096, 2048, 128)
     temperature = st.slider("Temperature", 0.0, 1.0, 0.1, 0.1) 
     top_p = st.slider("Top P", 0.1, 1.0, 0.95, 0.05)
@@ -107,44 +110,34 @@ with st.sidebar:
 
 # --- 初始化 Hugging Face LLM Client (重大替換) ---
 @st.cache_resource
-def load_huggingface_llm(model_id):
-    if AutoModelForCausalLM is None:
-        st.error("無法載入 Hugging Face 依賴，請安裝：pip install transformers torch accelerate bitsandbytes")
+def load_inference_client(model_id):
+    if not os.environ.get("HF_TOKEN"):
         return None
+        
     try:
-        # 使用量化 (4-bit) 減少記憶體消耗，這是運行 8B 模型的常見做法
-        tokenizer = AutoTokenizer.from_pretrained(model_id, trust_remote_code=True)
-        model = AutoModelForCausalLM.from_pretrained(
+        # 使用 InferenceClient 替換 AutoModelForCausalLM 的載入
+        client = InferenceClient(
             model_id,
-            torch_dtype=torch.bfloat16 if torch.cuda.is_available() else None,
-            device_map="auto", # <--- 讓 accelerate 管理裝置
-            trust_remote_code=True,
-            # load_in_4bit=True # 如果需要 4-bit 量化
-        )
-        # 使用 pipeline 簡化呼叫
-        llm_pipeline = pipeline(
-            "text-generation",
-            model=model,
-            tokenizer=tokenizer,
-            # device=(0 if torch.cuda.is_available() else -1) # <--- **移除此參數**
+            token=os.environ.get("HF_TOKEN")
         )
-        st.success(f"Hugging Face 模型 **{model_id}** 載入成功。")
-        return llm_pipeline
+        st.success(f"Hugging Face Inference Client **{model_id}** 載入成功。")
+        return client
     except Exception as e:
-        st.error(f"Hugging Face 模型載入失敗: {e}")
+        st.error(f"Hugging Face Inference Client 載入失敗: {e}")
         return None
 
-# 在 main 區塊外初始化 pipeline
-llm_pipeline = None
-if AutoModelForCausalLM is not None:
-    with st.spinner(f"正在載入 LLM 模型: {MODEL_ID} (8B)... (可能需要數分鐘)"):
-        llm_pipeline = load_huggingface_llm(MODEL_ID)
+# 在 main 區塊外初始化 client
+inference_client = None
+if os.environ.get("HF_TOKEN"):
+    with st.spinner(f"正在連線到 Inference Client: {MODEL_ID}..."):
+        inference_client = load_inference_client(MODEL_ID)
 
-if llm_pipeline is None:
-    st.warning("Hugging Face LLM 無法載入。請檢查依賴和環境資源。")
+if inference_client is None and os.environ.get("HF_TOKEN"):
+    st.warning("Hugging Face Inference Client 無法連線。請檢查您的 HF Token 和模型存取權限。")
+elif not os.environ.get("HF_TOKEN"):
+     st.error("請在環境變數中設定 HF_TOKEN 以啟用 LLM。")
 # =======================================================================
 
-
 # === Embedding 模型 (用於 RAG 參考庫) (保持不變) ===
 @st.cache_resource
 def load_embedding_model():
@@ -167,6 +160,7 @@ with st.spinner("正在載入 Embedding 模型..."):
 
 # === 建立向量庫 / Search 函數 (保持不變) ===
 def process_file_to_faiss(uploaded_file):
+    # 函數內容保持不變 (與原代碼相同)
     text_content = ""
     try:
         if uploaded_file.type == "application/pdf":
@@ -179,32 +173,32 @@ def process_file_to_faiss(uploaded_file):
         else:
             stringio = io.StringIO(uploaded_file.getvalue().decode("utf-8"))
             text_content = stringio.read()
-            
+                    
         if not text_content.strip():
             return None, "File is empty"
-        
+                
         # 嘗試以 </Event> 分割 Log，否則以換行符分割
         events = [e + "</Event>" for e in text_content.split("</Event>") if e.strip()]
         if len(events) <= 1:
             events = [line for line in text_content.split("\n") if line.strip()]
-            
+                    
         docs = [Document(page_content=e) for e in events]
-        
+                
         if not docs:
             return None, "No documents created"
-        
+                
         embeddings = embedding_model.embed_documents([d.page_content for d in docs])
         embeddings_np = np.array(embeddings).astype("float32")
         faiss.normalize_L2(embeddings_np) # L2 正規化
-        
+                
         dimension = embeddings_np.shape[1]
         index = faiss.IndexFlatIP(dimension) # IndexFlatIP (內積)
         index.add(embeddings_np)
-        
+                
         doc_ids = [str(uuid.uuid4()) for _ in range(len(docs))]
         docstore = InMemoryDocstore({_id: doc for _id, doc in zip(doc_ids, docs)})
         index_to_docstore_id = {i: _id for i, _id in enumerate(doc_ids)}
-        
+                
         vector_store = FAISS(
             embedding_function=embedding_model,
             index=index,
@@ -212,19 +206,20 @@ def process_file_to_faiss(uploaded_file):
             index_to_docstore_id=index_to_docstore_id,
             distance_strategy=DistanceStrategy.COSINE # 使用 Cosine 距離 (對應 IndexFlatIP)
         )
-        
+                
         return vector_store, f"{len(docs)} chunks created."
     except Exception as e:
         return None, f"Error: {str(e)}"
 
 def faiss_cosine_search_all(vector_store, query, threshold):
+    # 函數內容保持不變 (與原代碼相同)
     q_emb = embedding_model.embed_query(query)
     q_emb = np.array([q_emb]).astype("float32")
     faiss.normalize_L2(q_emb)
-    
+        
     index = vector_store.index
     D, I = index.search(q_emb, k=index.ntotal)
-    
+        
     selected = []
     for score, idx in zip(D[0], I[0]):
         if idx == -1: continue
@@ -233,20 +228,20 @@ def faiss_cosine_search_all(vector_store, query, threshold):
             doc_id = vector_store.index_to_docstore_id[idx]
             doc = vector_store.docstore.search(doc_id)
             selected.append((doc, score))
-            
+                
     selected.sort(key=lambda x: x[1], reverse=True)
     return selected
 
-# === Hugging Face 生成單一 Log 分析回答 (核心批量處理函數) (重大替換) ===
-def generate_rag_response_hf_for_log(llm_pipeline, model_id, log_sequence_text, user_prompt, sys_prompt, vector_store, threshold, max_output_tokens, temperature, top_p):
+# === Hugging Face 生成單一 Log 分析回答 (核心批量處理函數 - 重大替換為 InferenceClient) ===
+def generate_rag_response_hf_for_log(client, model_id, log_sequence_text, user_prompt, sys_prompt, vector_store, threshold, max_output_tokens, temperature, top_p):
     """
-    使用 Hugging Face LLM 執行 RAG 增強的 Log 序列分析。
+    使用 Hugging Face Inference Client 執行 RAG 增強的 Log 序列分析。
     """
-    if llm_pipeline is None:
-        return "ERROR: Hugging Face LLM Pipeline 未載入。", ""
-    
+    if client is None:
+        return "ERROR: Hugging Face Inference Client 未載入或 HF_TOKEN 未設定。", ""
+
     context_text = ""
-    # 1. RAG 檢索邏輯
+    # 1. RAG 檢索邏輯 (保持不變)
     if vector_store:
         selected = faiss_cosine_search_all(vector_store, log_sequence_text, threshold)
         if selected:
@@ -255,58 +250,53 @@ def generate_rag_response_hf_for_log(llm_pipeline, model_id, log_sequence_text,
                 for i, (doc, score) in enumerate(selected[:5]) # 限制檢索結果數量
             ]
             context_text = "\n".join(retrieved_contents)
-            
-    # 2. 建構 Prompt 的 RAG 部分和指令部分 (針對 HF 指令模型)
-    rag_instruction = f"""=== RETRIEVED REFERENCE CONTEXT (Cosine ≥ {threshold}) ===
-{context_text if context_text else 'No relevant reference context found.'}
-=== END REFERENCE CONTEXT ===
+
+    # 2. 建構 Llama 3 的 ChatML 格式的 Messages 列表
+    rag_instruction = f"""=== RETRIEVED REFERENCE CONTEXT (Cosine ≥ {threshold}) ==={context_text if context_text else 'No relevant reference context found.'}=== END REFERENCE CONTEXT ===
+    
 ANALYSIS INSTRUCTION: {user_prompt}
 Based on the provided LOG SEQUENCE and REFERENCE CONTEXT, you must analyze the **entire sequence** to detect any continuous attack chains or evolving threats. Focus on the **last log entry in the sequence** to determine its final criticality and priority, considering the preceding {WINDOW_SIZE} logs."""
-
-    log_content_section = f"""=== CURRENT LOG SEQUENCE TO ANALYZE (Window Size: {WINDOW_SIZE}) ===
-{log_sequence_text}
-=== END LOG SEQUENCE ==="""
     
-    # 整合 System Prompt、RAG、和 Log 內容
-    # 注意：fdtn-ai/Foundation-Sec-1.1-8B-Instruct 遵循 ChatML 格式，但此處使用簡化的 instruction-tuning 格式
-    full_prompt = (
-        f"**SYSTEM INSTRUCTION**: {sys_prompt}\n\n"
-        f"**RAG & ANALYSIS INSTRUCTION**:\n{rag_instruction}\n\n"
-        f"**LOG DATA**:\n{log_content_section}\n\n"
-        f"**RESPONSE**:"
-    )
+    log_content_section = f"""=== CURRENT LOG SEQUENCE TO ANALYZE (Window Size: {WINDOW_SIZE}) ===\n{log_sequence_text}\n=== END LOG SEQUENCE ==="""
     
-    # 3. 呼叫 Hugging Face Pipeline
+    # 整合 System Prompt、RAG、和 Log 內容到 messages 列表
+    # Llama 3 標準的 chat 格式
+    messages = [
+        {"role": "system", "content": sys_prompt},
+        {"role": "user", "content": f"{rag_instruction}\n\n{log_content_section}"}
+    ]
+
+    # 3. 呼叫 Hugging Face Inference Client
     try:
-        # Pipeline 參數設定
-        response = llm_pipeline(
-            full_prompt,
-            max_new_tokens=max_output_tokens,
+        # 使用 client.chat_completion 替換 pipeline 呼叫
+        response_stream = client.chat_completion(
+            messages,
+            max_tokens=max_output_tokens,
             temperature=temperature,
             top_p=top_p,
-            do_sample=True, # 啟用採樣
-            return_full_text=False # 只返回生成的文本
+            stream=False, # 由於是批量分析，不啟用流式輸出，一次性獲得結果
         )
-        
-        # 處理 pipeline 的輸出格式
-        if response and isinstance(response, list) and 'generated_text' in response[0]:
-            return response[0]['generated_text'].strip(), context_text
-        else:
-            return f"Hugging Face Pipeline 輸出格式錯誤: {response}", context_text
 
+        # 處理 chat_completion 的輸出格式 (非流式)
+        if response_stream and response_stream.choices:
+            # chat_completion 在非流式下返回一個 ChatCompletionResponse
+            generated_text = response_stream.choices[0].message.content
+            return generated_text.strip(), context_text
+        else:
+            return f"Hugging Face Inference Client 輸出格式錯誤: {response_stream}", context_text
+            
     except Exception as e:
         # 如果模型呼叫失敗，回傳詳細錯誤訊息
         return f"Hugging Face Model Error: {str(e)}", context_text
 
-
 # === 檔案處理和主執行邏輯 (保持結構，替換 LLM 呼叫) ===
-# 初始化 Session State
+# 初始化 Session State (保持不變)
 if 'execute_batch_analysis' not in st.session_state:
     st.session_state.execute_batch_analysis = False
 if 'batch_results' not in st.session_state:
     st.session_state.batch_results = None
-    
-# --- 1. 處理 RAG 知識庫檔案 (rag_uploaded_file) ---
+
+# --- 1. 處理 RAG 知識庫檔案 (rag_uploaded_file) --- (保持不變)
 if 'rag_current_file_key' not in st.session_state:
     st.session_state.rag_current_file_key = None
     
@@ -328,8 +318,8 @@ elif 'vector_store' in st.session_state:
     del st.session_state.vector_store
     del st.session_state.rag_current_file_key
     st.info("RAG 檔案已移除，已清除相關知識庫。")
-    
-# --- 2. 處理 JSON 批量分析檔案 (json_uploaded_file) ---
+
+# --- 2. 處理 JSON 批量分析檔案 (json_uploaded_file) --- (保持不變)
 if 'json_current_file_key' not in st.session_state:
     st.session_state.json_current_file_key = None
     
@@ -343,12 +333,11 @@ if json_uploaded_file:
             st.session_state.json_data_for_batch = json_data
             st.session_state.json_current_file_key = json_file_key
             st.toast("JSON Log 檔案已載入，請按 '執行批量分析'。", icon="📄")
-            
+                    
         except Exception as e:
             st.error(f"JSON 檔案解析錯誤: {e}")
             if 'json_data_for_batch' in st.session_state:
                 del st.session_state.json_data_for_batch
-                
 # 檔案移除/狀態清理 (如果使用者移除了 JSON 檔案)
 elif 'json_data_for_batch' in st.session_state:
     del st.session_state.json_data_for_batch
@@ -363,11 +352,10 @@ if st.session_state.execute_batch_analysis and 'json_data_for_batch' in st.sessi
     start_time = time.time() # 開始計時
     st.session_state.batch_results = []
     
-    if llm_pipeline is None:
-        st.error("Hugging Face LLM Pipeline 未載入，請檢查依賴和環境資源，無法執行批量分析。")
-        # 由於這是一個 Streamlit App，我們不直接 st.stop()，讓使用者可以檢查設定
+    if inference_client is None:
+        st.error("Hugging Face Inference Client 未載入，請檢查 HF_TOKEN 和網路連線，無法執行批量分析。")
         st.session_state.execute_batch_analysis = False
-    
+        
     data_to_process = st.session_state.json_data_for_batch
     
     # 提取 Log 列表的邏輯 (保持不變)
@@ -385,7 +373,7 @@ if st.session_state.execute_batch_analysis and 'json_data_for_batch' in st.sessi
             logs_list = [data_to_process]
     else:
         logs_list = [data_to_process]
-        
+    
     if logs_list:
         vs = st.session_state.get("vector_store", None)
         if vs:
@@ -393,8 +381,7 @@ if st.session_state.execute_batch_analysis and 'json_data_for_batch' in st.sessi
         else:
             st.warning("⚠️ RAG 知識庫未載入，將單純執行 Log 分析。")
             
-        # --- 新增：創建平移視窗序列 ---
-        
+        # --- 新增：創建平移視窗序列 --- (保持不變)
         # 將所有 Log 轉換為 JSON 格式化字串列表，以便後續拼接
         formatted_logs = [json.dumps(log, indent=2, ensure_ascii=False) for log in logs_list]
         
@@ -436,9 +423,9 @@ if st.session_state.execute_batch_analysis and 'json_data_for_batch' in st.sessi
             progress_bar.progress((i + 1) / total_sequences, text=f"已處理 {i + 1}/{total_sequences} 個序列 (目標 Log #{log_id})...")
             
             try:
-                # *** 替換為 Hugging Face 呼叫函數 ***
+                # *** 替換為 Inference Client 呼叫函數 ***
                 response, retrieved_ctx = generate_rag_response_hf_for_log(
-                    llm_pipeline=llm_pipeline, # <--- 新的 LLM pipeline
+                    client=inference_client, # <--- 新的 Inference Client
                     model_id=MODEL_ID,
                     log_sequence_text=seq_data["sequence_text"],
                     user_prompt=analysis_prompt,
@@ -460,7 +447,7 @@ if st.session_state.execute_batch_analysis and 'json_data_for_batch' in st.sessi
                 }
                 st.session_state.batch_results.append(item)
                 
-                # 結果顯示邏輯
+                # 結果顯示邏輯 (保持不變)
                 with results_container:
                     st.subheader(f"Log/Alert #{item['log_id']} (序列分析完成)")
                     with st.expander(f"序列內容 (包含 {len(seq_data['sequence_text'].split('--- Log Index'))-1} 條 Log)"):
@@ -519,7 +506,7 @@ if st.session_state.batch_results and not st.session_state.execute_batch_analysi
     for item in st.session_state.batch_results:
         log_content_str_for_report = json.dumps(item["log_content"], indent=2, ensure_ascii=False).replace("`", "\\`")
         full_report_chunks.append(f"---\n\n### Log/Alert #{item['log_id']}\n\n#### 原始內容\n```json\n{log_content_str_for_report}\n```\n\n#### LLM 分析結果\n{item['analysis_result']}\n")
-        
+    
     st.info(f"偵測到 {len(st.session_state.batch_results)} 條 Log 的歷史分析結果。")
     st.download_button(
         label="📥 下載上次的完整報告 (.md)",