Spaces:
Running
Running
from flask import Flask, request, jsonify | |
from transformers import AutoTokenizer | |
import os | |
app = Flask(__name__) | |
# 加载tokenizer | |
tokenizer = AutoTokenizer.from_pretrained(os.path.dirname(__file__)) | |
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 | |
def health(): | |
return jsonify({ | |
'status': 'healthy', | |
'tokenizer': 'deepseek-v3' | |
}) | |
if __name__ == '__main__': | |
app.run(host='127.0.0.1', port=7861) |