deepseek_v3_tokenizer/deepseek_service.py

from flask import Flask, request, jsonify
from transformers import AutoTokenizer
import os

app = Flask(__name__)

# 加载tokenizer
tokenizer = AutoTokenizer.from_pretrained(os.path.dirname(__file__))

@app.route('/count_tokens', methods=['POST'])
def count_tokens():
    try:
        data = request.json
        messages = data.get('messages', [])
        system = data.get('system')
        
        # 构建完整文本
        text = ""
        if system:
            text += f"System: {system}\n\n"
            
        for msg in messages:
            role = msg.get('role', '')
            content = msg.get('content', '')
            if role == 'user':
                text += f"User: {content}\n"
            elif role == 'assistant':
                text += f"Assistant: {content}\n"
            else:
                text += f"{role}: {content}\n"
        
        # 计算token数量
        tokens = tokenizer.encode(text)
        token_count = len(tokens)
        
        return jsonify({
            'input_tokens': token_count
        })
    except Exception as e:
        return jsonify({
            'error': str(e)
        }), 400

@app.route('/health', methods=['GET'])
def health():
    return jsonify({
        'status': 'healthy',
        'tokenizer': 'deepseek-v3'
    })

if __name__ == '__main__':
    app.run(host='127.0.0.1', port=7861)