pip install torch transformers flask

from flask import Flask, request, jsonify
from transformers import AutoModelForCausalLM, AutoTokenizer
import torch

app = Flask(__name__)

# 加載模型和分詞器
MODEL_NAME = "deepseek-ai/DeepSeek-R1-Distill-Qwen-1.5B"
tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)
model = AutoModelForCausalLM.from_pretrained(MODEL_NAME)

# 設(shè)置 pad_token_id
if tokenizer.pad_token isNone:
    tokenizer.pad_token = tokenizer.eos_token
model.config.pad_token_id = tokenizer.pad_token_id

# 設(shè)置設(shè)備
device = torch.device("cuda"if torch.cuda.is_available() else"cpu")
model.to(device)

def generate_text(prompt, max_length=100, top_p=0.95, temperature=0.7):
    """
    根據(jù)輸入的提示（prompt）生成文本。

    參數(shù):
        prompt (str): 輸入的文本提示，用于引導(dǎo)模型生成后續(xù)內(nèi)容。
        max_length (int): 生成文本的最大長(zhǎng)度（包括輸入提示）。
        top_p (float): 核采樣（Nucleus Sampling）的概率閾值，控制生成文本的多樣性。
        temperature (float): 溫度參數(shù)，控制生成文本的隨機(jī)性。

    返回:
        str: 生成的文本。
    """
    # 使用分詞器將輸入文本（prompt）轉(zhuǎn)換為模型所需的輸入格式。
    # return_tensors="pt" 表示返回 PyTorch 張量格式。
    # padding=True 和 truncatinotallow=True 確保輸入長(zhǎng)度一致且不超過(guò)模型的最大長(zhǎng)度。
    inputs = tokenizer(prompt, return_tensors="pt", padding=True, truncatinotallow=True).to(device)

    # 調(diào)用模型的 generate 方法生成文本。
    # inputs.input_ids 是輸入文本的編碼，attention_mask 是注意力掩碼，用于指示哪些位置是有效的輸入。
    outputs = model.generate(
        inputs.input_ids,  # 輸入文本的編碼
        attention_mask=inputs.attention_mask,  # 注意力掩碼
        max_length=max_length,  # 生成文本的最大長(zhǎng)度
        num_return_sequences=1,  # 返回的序列數(shù)量（這里只生成一個(gè)序列）
        no_repeat_ngram_size=2,  # 避免重復(fù)的 n-gram 大?。ǚ乐怪貜?fù)生成相同的短語(yǔ)）
        do_sample=True,  # 是否使用采樣（而不是貪婪解碼）
        top_p=top_p,  # 核采樣概率閾值
        temperature=temperature,  # 溫度參數(shù)，控制隨機(jī)性
        pad_token_id=tokenizer.pad_token_id,  # 填充標(biāo)記的 ID
    )

    # 將生成的編碼（outputs）解碼為可讀的文本。
    # skip_special_tokens=True 表示跳過(guò)特殊標(biāo)記（如 [CLS]、[SEP] 等）。
    generated_text = tokenizer.decode(outputs[0], skip_special_tokens=True)

    # 返回生成的文本。
    return generated_text

# API 路由
@app.route("/generate", methods=["POST"])
def generate():
    data = request.json
    prompt = data.get("prompt", "")
    max_length = data.get("max_length", 100)
    top_p = data.get("top_p", 0.95)
    temperature = data.get("temperature", 0.7)

    ifnot prompt:
        return jsonify({"error": "Prompt is required"}), 400

    try:
        generated_text = generate_text(prompt, max_length, top_p, temperature)
        return jsonify({"generated_text": generated_text})
    except Exception as e:
        return jsonify({"error": str(e)}), 500

# 啟動(dòng)應(yīng)用
if __name__ == "__main__":
    app.run(host="0.0.0.0", port=5000)

python app.py

set HF_ENDPOINT=https://hf-mirror.com

set HF_ENDPOINT=https://hf-mirror.com

export HF_ENDPOINT=https://hf-mirror.com

export HF_ENDPOINT="https://hf-mirror.com"