pip install torch transformers langchain  # 直接裝到全局

# 方式一：使用 venv（推薦，Python 內(nèi)置）
python -m venv ai-qa-env
source ai-qa-env/bin/activate        # Linux/macOS
ai-qa-env\Scripts\activate           # Windows
# 方式二：使用 conda
conda create -n ai-qa python=3.11
conda activate ai-qa

nvidia-smi  # 查看 GPU 驅(qū)動支持的最高 CUDA 版本
nvcc --version  # 查看已安裝的 CUDA Toolkit 版本

# CUDA 12.1 對應(yīng)的安裝命令示例
pip install torch torchvision torchaudio --index-url https://download.pytorch.org/whl/cu121
# 沒有 GPU，只用 CPU
pip install torch torchvision torchaudio --index-url https://download.pytorch.org/whl/cpu

import torch

print(f"PyTorch 版本: {torch.__version__}")
print(f"CUDA 是否可用: {torch.cuda.is_available()}")
print(f"GPU 數(shù)量: {torch.cuda.device_count()}")

if torch.cuda.is_available():
    print(f"當(dāng)前 GPU: {torch.cuda.get_device_name(0)}")
    print(f"顯存總量: {torch.cuda.get_device_properties(0).total_memory / 1024**3:.1f} GB")

# 生成當(dāng)前環(huán)境的依賴快照
pip freeze > requirements.txt
# 在新環(huán)境中還原
pip install -r requirements.txt

torch==2.2.0
transformers==4.38.0
langchain==0.1.9
langchain-community==0.0.24
faiss-cpu==1.7.4
sentence-transformers==2.5.1
ollama==0.1.7
gradio==4.19.2

# 安裝 Ollama（macOS/Linux）
curl -fsSL https://ollama.com/install.sh | sh
# 下載并運(yùn)行模型
ollama pull llama3.2        # Meta Llama 3.2 (3B)
ollama pull qwen2.5:7b      # 阿里通義千問 2.5 (7B)
ollama pull deepseek-r1:7b  # DeepSeek R1 (7B)
# 驗(yàn)證模型列表
ollama list

import ollama

def ask_local_llm(question: str, model: str = "qwen2.5:7b") -> str:
    """
    調(diào)用本地 Ollama 模型進(jìn)行問答
    
    Args:
        question: 用戶問題
        model: 模型名稱
    
    Returns:
        模型回答
    """
    response = ollama.chat(
        model=model,
        messages=[
            {
                "role": "system",
                "content": "你是一個(gè)專業(yè)的 AI 助手，請用中文簡潔準(zhǔn)確地回答問題。"
            },
            {
                "role": "user", 
                "content": question
            }
        ]
    )
    return response["message"]["content"]


# 測試
if __name__ == "__main__":
    answer = ask_local_llm("Python 中的 GIL 是什么？")
    print(answer)

"""
本地 RAG 問答系統(tǒng)
依賴: pip install langchain langchain-community faiss-cpu sentence-transformers ollama
"""

from langchain.text_splitter import RecursiveCharacterTextSplitter
from langchain_community.vectorstores import FAISS
from langchain_community.embeddings import HuggingFaceEmbeddings
from langchain_community.document_loaders import TextLoader, DirectoryLoader
from langchain.chains import RetrievalQA
from langchain_community.llms import Ollama
import os


class LocalRAGSystem:
    """本地 RAG 問答系統(tǒng)"""

    def __init__(
        self,
        docs_dir: str = "./docs",
        model_name: str = "qwen2.5:7b",
        embedding_model: str = "BAAI/bge-small-zh-v1.5",
        chunk_size: int = 500,
        chunk_overlap: int = 50,
    ):
        self.docs_dir = docs_dir
        self.model_name = model_name
        self.chunk_size = chunk_size
        self.chunk_overlap = chunk_overlap

        print("?? 初始化 Embedding 模型...")
        # 使用本地 Embedding 模型，避免調(diào)用外部 API
        self.embeddings = HuggingFaceEmbeddings(
            model_name=embedding_model,
            model_kwargs={"device": "cpu"},  # 改為 "cuda" 可用 GPU 加速
            encode_kwargs={"normalize_embeddings": True},
        )

        self.vectorstore = None
        self.qa_chain = None

    def load_and_index(self):
        """加載文檔并建立向量索引"""
        print(f"?? 加載文檔目錄: {self.docs_dir}")

        # 支持多種文檔格式
        loader = DirectoryLoader(
            self.docs_dir,
            glob="**/*.txt",
            loader_cls=TextLoader,
            loader_kwargs={"encoding": "utf-8"},
        )
        documents = loader.load()
        print(f"? 加載了 {len(documents)} 個(gè)文檔")

        # 文本分塊
        splitter = RecursiveCharacterTextSplitter(
            chunk_size=self.chunk_size,
            chunk_overlap=self.chunk_overlap,
            separators=["\n\n", "\n", "。", "！", "？", " ", ""],
        )
        chunks = splitter.split_documents(documents)
        print(f"? 分割為 {len(chunks)} 個(gè)文本塊")

        # 建立向量索引
        print("?? 建立向量索引（首次較慢，請耐心等待）...")
        self.vectorstore = FAISS.from_documents(chunks, self.embeddings)
        print("? 向量索引建立完成")

        # 保存索引到本地（下次直接加載，無需重建）
        self.vectorstore.save_local("./faiss_index")
        print("?? 索引已保存到 ./faiss_index")

    def load_existing_index(self):
        """加載已有的向量索引"""
        if os.path.exists("./faiss_index"):
            print("?? 加載已有向量索引...")
            self.vectorstore = FAISS.load_local(
                "./faiss_index",
                self.embeddings,
                allow_dangerous_deserialization=True,
            )
            print("? 索引加載完成")
        else:
            print("??  未找到已有索引，請先調(diào)用 load_and_index()")

    def build_qa_chain(self):
        """構(gòu)建問答鏈"""
        if self.vectorstore is None:
            raise ValueError("請先調(diào)用 load_and_index() 或 load_existing_index()")

        print(f"?? 連接本地 LLM: {self.model_name}")
        llm = Ollama(
            model=self.model_name,
            temperature=0.1,  # 降低隨機(jī)性，讓回答更穩(wěn)定
        )

        retriever = self.vectorstore.as_retriever(
            search_type="similarity",
            search_kwargs={"k": 3},  # 召回最相關(guān)的 3 個(gè)文檔塊
        )

        self.qa_chain = RetrievalQA.from_chain_type(
            llm=llm,
            chain_type="stuff",
            retriever=retriever,
            return_source_documents=True,
        )
        print("? 問答系統(tǒng)就緒！")

    def ask(self, question: str) -> dict:
        """
        提問并獲取答案
        
        Returns:
            dict: {"answer": str, "sources": list}
        """
        if self.qa_chain is None:
            raise ValueError("請先調(diào)用 build_qa_chain()")

        result = self.qa_chain.invoke({"query": question})

        return {
            "answer": result["result"],
            "sources": [
                doc.metadata.get("source", "未知來源")
                for doc in result["source_documents"]
            ],
        }


# ============ 使用示例 ============
if __name__ == "__main__":
    # 初始化系統(tǒng)
    rag = LocalRAGSystem(
        docs_dir="./my_docs",
        model_name="qwen2.5:7b",
    )

    # 首次使用：加載文檔并建立索引
    rag.load_and_index()

    # 后續(xù)使用：直接加載已有索引（更快）
    # rag.load_existing_index()

    # 構(gòu)建問答鏈
    rag.build_qa_chain()

    # 開始問答
    while True:
        question = input("\n? 請輸入問題（輸入 q 退出）: ").strip()
        if question.lower() == "q":
            break

        result = rag.ask(question)
        print(f"\n?? 回答:\n{result['answer']}")
        print(f"\n?? 參考來源: {', '.join(result['sources'])}")

from transformers import AutoModelForCausalLM, AutoTokenizer, BitsAndBytesConfig
import torch

# 4-bit 量化配置（顯存減少約 75%）
quantization_config = BitsAndBytesConfig(
    load_in_4bit=True,
    bnb_4bit_compute_dtype=torch.float16,
    bnb_4bit_use_double_quant=True,
    bnb_4bit_quant_type="nf4",
)

model_id = "Qwen/Qwen2.5-7B-Instruct"

tokenizer = AutoTokenizer.from_pretrained(model_id)
model = AutoModelForCausalLM.from_pretrained(
    model_id,
    quantization_config=quantization_config,
    device_map="auto",  # 自動分配到 GPU/CPU
)

print(f"模型加載完成，占用顯存: {torch.cuda.memory_allocated() / 1024**3:.2f} GB")

import ollama

def stream_answer(question: str, model: str = "qwen2.5:7b"):
    """流式輸出，邊生成邊顯示"""
    print("?? ", end="", flush=True)
    
    for chunk in ollama.chat(
        model=model,
        messages=[{"role": "user", "content": question}],
        stream=True,  # 開啟流式輸出
    ):
        content = chunk["message"]["content"]
        print(content, end="", flush=True)
    
    print()  # 換行


stream_answer("用一句話解釋什么是 Transformer 架構(gòu)")

import gradio as gr
from local_rag import LocalRAGSystem  # 引用上面的代碼

# 初始化 RAG 系統(tǒng)
rag = LocalRAGSystem()
rag.load_existing_index()
rag.build_qa_chain()


def chat(message: str, history: list) -> str:
    """Gradio 聊天回調(diào)函數(shù)"""
    if not message.strip():
        return "請輸入問題"
    
    result = rag.ask(message)
    answer = result["answer"]
    sources = result["sources"]
    
    if sources:
        answer += f"\n\n---\n?? **參考來源**: {', '.join(set(sources))}"
    
    return answer


# 創(chuàng)建 Gradio 界面
demo = gr.ChatInterface(
    fn=chat,
    title="?? 本地 AI 問答系統(tǒng)",
    description="基于本地 LLM + RAG 的私有知識庫問答",
    examples=[
        "這個(gè)系統(tǒng)是如何工作的？",
        "請總結(jié)一下主要內(nèi)容",
    ],
    theme=gr.themes.Soft(),
)

if __name__ == "__main__":
    demo.launch(
        server_name="0.0.0.0",
        server_port=7860,
        share=False,  # 改為 True 可生成公網(wǎng)鏈接
    )

# 設(shè)置 HuggingFace 鏡像（國內(nèi)訪問加速）
export HF_ENDPOINT=https://hf-mirror.com
# pip 使用清華鏡像
pip install -i https://pypi.tuna.tsinghua.edu.cn/simple transformers
# conda 使用清華鏡像
conda config --add channels https://mirrors.tuna.tsinghua.edu.cn/anaconda/pkgs/main

# 在代碼中指定鏡像
import os
os.environ["HF_ENDPOINT"] = "https://hf-mirror.com"

from transformers import AutoTokenizer
# 之后的下載會自動走鏡像
tokenizer = AutoTokenizer.from_pretrained("Qwen/Qwen2.5-7B-Instruct")