AI_Gemma实践笔记

# 操作系统
Ubuntu 22.04 LTS

# Python 环境
Python 3.10.12
conda 环境

# 硬件
CPU: Intel i7-12700K
RAM: 16GB DDR5
GPU: NVIDIA RTX 3060 12GB

# 创建 conda 环境
conda create -n gemma python=3.10 -y
conda activate gemma

# 安装 PyTorch（带 CUDA 支持）
pip install torch torchvision torchaudio --index-url https://download.pytorch.org/whl/cu118

# 安装 transformers 和相关库
pip install transformers accelerate bitsandbytes sentencepiece

pip install huggingface_hub
huggingface-cli login
# 输入你的 token

from transformers import AutoTokenizer, AutoModelForCausalLM
import torch

model_id = "google/gemma-2b"

# 加载 tokenizer
tokenizer = AutoTokenizer.from_pretrained(model_id)

# 加载模型（使用 4-bit 量化节省显存）
model = AutoModelForCausalLM.from_pretrained(
    model_id,
    device_map="auto",
    load_in_4bit=True,
    torch_dtype=torch.float16
)

export HF_ENDPOINT=https://hf-mirror.com
# 然后再下载就快多了

def generate_response(prompt, max_length=512):
    input_ids = tokenizer(prompt, return_tensors="pt").input_ids.to("cuda")

    with torch.no_grad():
        outputs = model.generate(
            input_ids,
            max_length=max_length,
            do_sample=True,
            top_k=50,
            top_p=0.95,
            temperature=0.7,
            repetition_penalty=1.1
        )

    response = tokenizer.decode(outputs[0], skip_special_tokens=True)
    return response

# 测试
question = "什么是机器学习？"
answer = generate_response(question)
print(answer)

pip install vllm

from vllm import LLM, SamplingParams

llm = LLM(model="google/gemma-2b", trust_remote_code=True)
sampling_params = SamplingParams(temperature=0.7, top_p=0.95, max_tokens=512)

outputs = llm.generate(["什么是机器学习？"], sampling_params)
print(outputs[0].outputs[0].text)

# 安装转换工具
pip install llama-cpp-python[server]

# 下载转换脚本
wget https://raw.githubusercontent.com/ggerganov/llama.cpp/master/convert-hf-to-gguf.py

# 转换模型
python convert-hf-to-gguf.py /path/to/gemma-2b --outfile gemma-2b.gguf

from llama_cpp import Llama

llm = Llama(
    model_path="gemma-2b.gguf",
    n_ctx=4096,
    n_gpu_layers=-1,  # 使用所有 GPU 层
    verbose=False
)

output = llm("什么是机器学习？", max_tokens=512)
print(output['choices'][0]['text'])

class ChatSession:
    def __init__(self, model, tokenizer):
        self.model = model
        self.tokenizer = tokenizer
        self.past_key_values = None
        self.history = []

    def chat(self, message):
        self.history.append({"role": "user", "content": message})
        prompt = self.tokenizer.apply_chat_template(
            self.history,
            tokenize=False,
            add_generation_prompt=True
        )

        inputs = self.tokenizer(prompt, return_tensors="pt").to("cuda")

        with torch.no_grad():
            outputs = self.model.generate(
                **inputs,
                past_key_values=self.past_key_values,
                max_length=inputs.input_ids.shape[1] + 512,
                do_sample=True,
                temperature=0.7,
                top_p=0.95
            )

        response = tokenizer.decode(outputs[0][inputs.input_ids.shape[1]:], skip_special_tokens=True)
        self.history.append({"role": "assistant", "content": response})

        # 保存 KV cache 用于下一轮
        self.past_key_values = outputs.past_key_values

        return response

model = AutoModelForCausalLM.from_pretrained(
    model_id,
    load_in_4bit=True,  # 关键！
    device_map="auto"
)

# 8-bit 量化
model = AutoModelForCausalLM.from_pretrained(
    model_id,
    load_in_8bit=True,
    device_map="auto"
)

# 或者纯 CPU
model = AutoModelForCausalLM.from_pretrained(
    model_id,
    device_map="cpu"
)

tokenizer = AutoTokenizer.from_pretrained(
    model_id,
    use_fast=True,
    trust_remote_code=True
)

# 或者用专门训练中文的变体
model_id = "THUDM/chatglm3-6b"  # 这个中文效果更好

outputs = model.generate(
    input_ids,
    repetition_penalty=1.1,  # 防止重复
    # 其他参数...
)

# 使用支持长文本的版本
model_id = "google/gemma-2b-it"  # "it" 表示 instruction-tuned 版本

# 或者增加 context window
model = AutoModelForCausalLM.from_pretrained(
    model_id,
    max_position_embeddings=8192,  # 扩展到 8K
    # ...
)

pip install vllm==0.2.5
pip install flash-attn==2.3.3

import time

def benchmark(model, test_questions):
    results = []
    for question in test_questions:
        start = time.time()
        answer = generate_response(question)
        end = time.time()
        results.append({
            "question": question,
            "time": end - start,
            "tokens": len(answer)
        })
    return results

test_questions = [
    "什么是机器学习？",
    "解释一下神经网络的工作原理",
    "如何优化深度学习模型的性能？",
    "什么是 Transformer 架构？",
    "对比一下 CNN 和 RNN 的优缺点"
]

from langchain.embeddings import HuggingFaceEmbeddings
from langchain.vectorstores import Chroma
from langchain.chains import RetrievalQA

# 加载文档
embeddings = HuggingFaceEmbeddings(model_name="BAAI/bge-base-zh")
vectorstore = Chroma.from_documents(documents, embeddings)

# 创建问答链
qa_chain = RetrievalQA.from_chain_type(
    llm=gemma_llm,
    chain_type="stuff",
    retriever=vectorstore.as_retriever()
)

# 查询
answer = qa_chain.run("我在哪篇文章里写过关于多线程的内容？")

code_prompt = """
用 Python 写一个函数，判断一个字符串是否是回文。
要考虑大小写和空格。
"""

code = generate_response(code_prompt, max_length=256)
print(code)

def summarize(text, max_length=200):
    prompt = f"请用中文总结以下内容的要点：\n\n{text}\n\n总结："
    summary = generate_response(prompt, max_length=max_length)
    return summary