transformers>=4.35.0
sentence-transformers>=2.2.2
faiss-cpu>=1.7.4
langchain>=0.1.0
langchain-community>=0.0.10
torch>=2.0.0
accelerate>=0.25.0
pypdf>=3.17.0
python-dotenv>=1.0.0

pip install -r requirements.txt

# knowledge_base.py
import os
from typing import List, Dict
import re
class KnowledgeBase:
    """本地知識庫管理類"""
    def __init__(self, chunk_size: int = 500, chunk_overlap: int = 50):
        """
        初始化知識庫
        Args:
            chunk_size: 文本塊的最大字符數(shù)
            chunk_overlap: 相鄰文本塊之間的重疊字符數(shù)
        """
        self.chunk_size = chunk_size
        self.chunk_overlap = chunk_overlap
        self.documents = []  # 存儲所有文本塊
    def load_markdown_file(self, file_path: str) -> str:
        """
        加載 Markdown 文件內(nèi)容
        Args:
            file_path: 文件路徑
        Returns:
            文件文本內(nèi)容
        """
        with open(file_path, 'r', encoding='utf-8') as f:
            return f.read()
    def split_text_into_chunks(self, text: str, source: str = "") -> List[Dict]:
        """
        將長文本切分為小塊
        Args:
            text: 待切分的文本
            source: 文本來源標(biāo)識（如文件名）
        Returns:
            文本塊列表，每個塊包含內(nèi)容和元數(shù)據(jù)
        """
        chunks = []
        start = 0
        text_length = len(text)
        while start < text_length:
            end = start + self.chunk_size
            # 如果不是最后一塊，嘗試在句號、問號或換行處切分
            if end < text_length:
                # 尋找最近的句子結(jié)束符
                for delimiter in ['。', '！', '？', '\n\n', '. ', '! ', '? ']:
                    pos = text.rfind(delimiter, start, end)
                    if pos != -1:
                        end = pos + len(delimiter)
                        break
            chunk_text = text[start:end].strip()
            if chunk_text:  # 跳過空塊
                chunks.append({
                    'content': chunk_text,
                    'metadata': {
                        'source': source,
                        'start': start,
                        'end': end
                    }
                })
            # 移動到下一塊，考慮重疊區(qū)域
            start = end - self.chunk_overlap
        return chunks
    def load_directory(self, directory: str, extension: str = '.md') -> None:
        """
        加載目錄下所有指定擴展名的文件
        Args:
            directory: 目錄路徑
            extension: 文件擴展名
        """
        for filename in os.listdir(directory):
            if filename.endswith(extension):
                file_path = os.path.join(directory, filename)
                print(f"正在加載文件: {filename}")
                content = self.load_markdown_file(file_path)
                chunks = self.split_text_into_chunks(content, source=filename)
                self.documents.extend(chunks)
        print(f"知識庫構(gòu)建完成，共 {len(self.documents)} 個文本塊")
    def get_documents(self) -> List[Dict]:
        """獲取所有文檔塊"""
        return self.documents

# vector_store.py
import numpy as np
import faiss
from sentence_transformers import SentenceTransformer
from typing import List, Dict, Tuple

class VectorStore:
    """向量存儲和檢索類"""

    def __init__(self, model_name: str = 'sentence-transformers/all-MiniLM-L6-v2'):
        """
        初始化向量存儲

        Args:
            model_name: 嵌入模型名稱
        """
        print(f"正在加載嵌入模型: {model_name}")
        self.embedding_model = SentenceTransformer(model_name)
        self.dimension = self.embedding_model.get_sentence_embedding_dimension()
        self.index = None
        self.documents = []  # 存儲原始文檔塊

    def build_index(self, documents: List[Dict]) -> None:
        """
        為文檔構(gòu)建 FAISS 索引

        Args:
            documents: 文檔塊列表
        """
        self.documents = documents
        print(f"正在為 {len(documents)} 個文檔塊生成嵌入向量...")

        # 批量生成嵌入向量
        texts = [doc['content'] for doc in documents]
        embeddings = self.embedding_model.encode(
            texts,
            show_progress_bar=True,
            convert_to_numpy=True
        )

        # 構(gòu)建 FAISS 索引（使用 L2 距離）
        self.index = faiss.IndexFlatL2(self.dimension)
        self.index.add(embeddings.astype('float32'))

        print(f"索引構(gòu)建完成，向量維度: {self.dimension}, 索引大小: {self.index.ntotal}")

    def save_index(self, index_path: str, docs_path: str) -> None:
        """
        保存索引和文檔到磁盤

        Args:
            index_path: FAISS 索引保存路徑
            docs_path: 文檔保存路徑
        """
        if self.index is None:
            raise ValueError("索引尚未構(gòu)建")

        faiss.write_index(self.index, index_path)
        np.save(docs_path, self.documents)
        print(f"索引已保存到 {index_path}")
        print(f"文檔已保存到 {docs_path}")

    def load_index(self, index_path: str, docs_path: str) -> None:
        """
        從磁盤加載索引和文檔

        Args:
            index_path: FAISS 索引路徑
            docs_path: 文檔路徑
        """
        self.index = faiss.read_index(index_path)
        self.documents = np.load(docs_path, allow_pickle=True).tolist()
        print(f"索引已加載，包含 {self.index.ntotal} 個向量")

    def search(self, query: str, top_k: int = 3) -> List[Dict]:
        """
        搜索與查詢最相似的文檔

        Args:
            query: 查詢文本
            top_k: 返回最相似的前 k 個結(jié)果

        Returns:
            包含文檔內(nèi)容和相似度分?jǐn)?shù)的列表
        """
        if self.index is None:
            raise ValueError("索引尚未構(gòu)建或加載")

        # 將查詢轉(zhuǎn)換為向量
        query_vector = self.embedding_model.encode(
            [query],
            convert_to_numpy=True
        ).astype('float32')

        # 搜索最相似的向量
        distances, indices = self.index.search(query_vector, top_k)

        # 構(gòu)建結(jié)果列表
        results = []
        for i, (distance, idx) in enumerate(zip(distances[0], indices[0])):
            if idx < len(self.documents):  # 確保索引有效
                result = self.documents[idx].copy()
                # 將 L2 距離轉(zhuǎn)換為相似度分?jǐn)?shù)（0-1 之間）
                result['score'] = 1 / (1 + distance)
                result['rank'] = i + 1
                results.append(result)

        return results

# llm_generator.py
import torch
from transformers import AutoTokenizer, AutoModelForCausalLM
from typing import List, Dict

class LLMGenerator:
    """LLM 答案生成器"""

    def __init__(
        self,
        model_name: str = 'microsoft/Phi-3-mini-4k-instruct',
        device: str = 'auto'
    ):
        """
        初始化 LLM 生成器

        Args:
            model_name: 模型名稱或本地路徑
            device: 運行設(shè)備 ('cuda', 'cpu', 'auto')
        """
        if device == 'auto':
            self.device = 'cuda' if torch.cuda.is_available() else 'cpu'
        else:
            self.device = device

        print(f"正在加載模型: {model_name}")
        print(f"使用設(shè)備: {self.device}")

        self.tokenizer = AutoTokenizer.from_pretrained(
            model_name,
            trust_remote_code=True
        )
        self.model = AutoModelForCausalLM.from_pretrained(
            model_name,
            torch_dtype=torch.float16 if self.device == 'cuda' else torch.float32,
            device_map='auto' if self.device == 'cuda' else None,
            trust_remote_code=True
        )

        if self.device == 'cpu':
            self.model = self.model.to(self.device)

        print("模型加載完成")

    def generate_answer(
        self,
        query: str,
        context_docs: List[Dict],
        max_new_tokens: int = 256,
        temperature: float = 0.7
    ) -> str:
        """
        基于上下文生成答案

        Args:
            query: 用戶問題
            context_docs: 檢索到的相關(guān)文檔
            max_new_tokens: 最大生成 token 數(shù)
            temperature: 采樣溫度（越低越確定）

        Returns:
            生成的答案
        """
        # 構(gòu)建上下文
        context_text = "\n\n".join([
            f"[參考信息 {i+1}]\n{doc['content']}"
            for i, doc in enumerate(context_docs)
        ])

        # 構(gòu)建提示詞（Prompt Engineering）
        prompt = f"""你是一個專業(yè)的智能客服助手。請根據(jù)以下參考信息回答用戶的問題。

參考信息：
{context_text}

用戶問題：{query}

請基于參考信息給出準(zhǔn)確、簡潔、友好的回答。如果參考信息中沒有相關(guān)內(nèi)容，請禮貌告知用戶你無法從現(xiàn)有資料中找到答案。

回答："""

        # 編碼輸入
        inputs = self.tokenizer(
            prompt,
            return_tensors='pt',
            truncation=True,
            max_length=2048  # 控制輸入長度
        ).to(self.device)

        # 生成答案
        with torch.no_grad():
            outputs = self.model.generate(
                **inputs,
                max_new_tokens=max_new_tokens,
                temperature=temperature,
                do_sample=True if temperature > 0 else False,
                top_p=0.9,
                repetition_penalty=1.1,
                pad_token_id=self.tokenizer.eos_token_id
            )

        # 解碼輸出（只返回新生成的部分）
        answer = self.tokenizer.decode(
            outputs[0][inputs['input_ids'].shape[1]:],
            skip_special_tokens=True
        )

        return answer.strip()

# main.py
import os
from dotenv import load_dotenv
from knowledge_base import KnowledgeBase
from vector_store import VectorStore
from llm_generator import LLMGenerator

class QASystem:
    """智能問答系統(tǒng)主類"""

    def __init__(
        self,
        knowledge_dir: str = './knowledge',
        index_path: str = './faiss_index.bin',
        docs_path: str = './documents.npy',
        rebuild_index: bool = False
    ):
        """
        初始化問答系統(tǒng)

        Args:
            knowledge_dir: 知識庫目錄
            index_path: FAISS 索引路徑
            docs_path: 文檔存儲路徑
            rebuild_index: 是否重建索引
        """
        # 初始化組件
        self.kb = KnowledgeBase(chunk_size=500, chunk_overlap=50)
        self.vector_store = VectorStore()
        self.llm = LLMGenerator()

        # 加載或構(gòu)建索引
        if rebuild_index or not os.path.exists(index_path):
            print("正在構(gòu)建新索引...")
            self.kb.load_directory(knowledge_dir)
            documents = self.kb.get_documents()
            self.vector_store.build_index(documents)
            self.vector_store.save_index(index_path, docs_path)
        else:
            print("正在加載已有索引...")
            self.vector_store.load_index(index_path, docs_path)

        print("=" * 50)
        print("智能問答系統(tǒng)初始化完成！")
        print("=" * 50)

    def ask(self, query: str, top_k: int = 3) -> Dict:
        """
        處理用戶問題

        Args:
            query: 用戶問題
            top_k: 檢索的文檔數(shù)量

        Returns:
            包含答案和參考信息的字典
        """
        print(f"\n用戶問題: {query}")
        print("正在檢索相關(guān)知識...")

        # 向量檢索
        retrieved_docs = self.vector_store.search(query, top_k=top_k)

        print(f"檢索到 {len(retrieved_docs)} 個相關(guān)文檔片段")

        # 生成答案
        print("正在生成答案...")
        answer = self.llm.generate_answer(query, retrieved_docs)

        return {
            'query': query,
            'answer': answer,
            'sources': [
                {
                    'content': doc['content'][:100] + '...',
                    'score': doc['score'],
                    'source': doc['metadata']['source']
                }
                for doc in retrieved_docs
            ]
        }

    def run_interactive(self):
        """運行交互式問答循環(huán)"""
        print("\n" + "=" * 50)
        print("智能問答系統(tǒng)已就緒！")
        print("輸入您的問題，輸入 'quit' 或 'exit' 退出")
        print("=" * 50 + "\n")

        while True:
            try:
                query = input("您的問題: ").strip()

                if not query:
                    continue

                if query.lower() in ['quit', 'exit', '退出']:
                    print("感謝使用，再見！")
                    break

                # 處理問題
                result = self.ask(query)

                # 顯示結(jié)果
                print(f"\n{'='*50}")
                print(f"答案: {result['answer']}")
                print(f"{'='*50}")
                print(f"參考信息 (共 {len(result['sources'])} 條):")
                for i, source in enumerate(result['sources'], 1):
                    print(f"\n{i}. 來源: {source['source']}")
                    print(f"   相似度: {source['score']:.3f}")
                    print(f"   內(nèi)容: {source['content']}")

            except KeyboardInterrupt:
                print("\n\n檢測到中斷，正在退出...")
                break
            except Exception as e:
                print(f"\n發(fā)生錯誤: {str(e)}")
                print("請嘗試重新提問\n")

if __name__ == '__main__':
    # 創(chuàng)建示例知識庫目錄和文件
    os.makedirs('./knowledge', exist_ok=True)

    # 創(chuàng)建示例知識庫文件（如果不存在）
    sample_md = """# 公司員工手冊

## 密碼重置流程

如果員工忘記了系統(tǒng)登錄密碼，可以按照以下步驟進(jìn)行重置：

1. 訪問公司內(nèi)部系統(tǒng)登錄頁面
2. 點擊"忘記密碼"鏈接
3. 輸入員工郵箱地址
4. 查收郵箱中的重置鏈接
5. 點擊鏈接進(jìn)入密碼重置頁面
6. 設(shè)置新密碼（需包含大小寫字母、數(shù)字和特殊字符）
7. 使用新密碼登錄系統(tǒng)

如遇到問題，請聯(lián)系 IT 部門熱線：400-123-4567

## 年假申請流程

員工申請年假需要遵循以下流程：

1. 提前至少 3 個工作日提交申請
2. 登錄 HR 系統(tǒng)，進(jìn)入"假期管理"模塊
3. 選擇年假類型，填寫起止日期
4. 填寫假期期間的工作交接安排
5. 提交申請等待審批
6. 直屬上級審批通過后生效

注意事項：
- 年假最小單位為 0.5 天
- 法定節(jié)假日前后 3 天內(nèi)不得申請年假
- 年假需在當(dāng)年 12 月 31 日前使用完畢

## 報銷流程

員工日常費用報銷流程如下：

1. 收集并整理所有原始發(fā)票
2. 登錄財務(wù)系統(tǒng)，進(jìn)入"費用報銷"模塊
3. 填寫報銷單，包括：
   - 費用類型（交通、餐飲、辦公用品等）
   - 金額和發(fā)票張數(shù)
   - 費用事由
4. 上傳發(fā)票掃描件或照片
5. 提交申請
6. 等待財務(wù)部門審核（一般 3-5 個工作日）
7. 審核通過后，費用將轉(zhuǎn)入員工工資卡

報銷標(biāo)準(zhǔn)：
- 市內(nèi)交通：單次不超過 50 元
- 餐飲費用：人均不超過 80 元/餐
- 住宿費用：一線城市不超過 500 元/晚

## 聯(lián)系方式

各部門聯(lián)系方式：
- IT 部門：it@company.com | 內(nèi)線 8888
- HR 部門：hr@company.com | 內(nèi)線 6666
- 財務(wù)部門：finance@company.com | 內(nèi)線 9999
"""

    sample_path = './knowledge/員工手冊.md'
    if not os.path.exists(sample_path):
        with open(sample_path, 'w', encoding='utf-8') as f:
            f.write(sample_md)
        print(f"已創(chuàng)建示例知識庫文件: {sample_path}")

    # 初始化并運行問答系統(tǒng)
    qa_system = QASystem(rebuild_index=True)
    qa_system.run_interactive()