File size: 1,656 Bytes
3b80f4a d04833c a0269f5 3b80f4a d04833c 3b80f4a cd6a397 3b80f4a fe0f77f 3b80f4a cd6a397 3b80f4a |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 |
from flask import Flask, request, jsonify
from sentence_transformers import SentenceTransformer
import logging
import os
from functools import wraps
logging.basicConfig(level=logging.DEBUG)
API_TOKEN = os.getenv("API_TOKEN", "your-default-token") # Set a default for local testing
app = Flask(__name__)
def require_token(f):
@wraps(f)
def decorated(*args, **kwargs):
auth_header = request.headers.get("Authorization")
if not auth_header:
return jsonify({"error": "Authorization header missing"}), 401
try:
token = auth_header.split("Bearer ")[1]
if token != API_TOKEN:
return jsonify({"error": "Invalid token"}), 403
except IndexError:
return jsonify({"error": "Invalid Authorization header format"}), 401
return f(*args, **kwargs)
return decorated
# Load the Qwen3-Embedding-0.6B model (adjust the path as needed)
model = SentenceTransformer(model_name_or_path="Qwen/Qwen3-Embedding-0.6B")
@app.route('/')
def home():
return 'Hello from Flask on Hugging Face!'
@app.route('/embed', methods=['POST'])
@require_token
def get_embedding():
try:
text = request.json.get('text', '')
if not text:
return jsonify({"error": "No text provided"}), 400
# Generate embeddings for the input text
document_embeddings = model.encode(text)
arr_list = document_embeddings.tolist()
return jsonify({"embedding": arr_list})
except Exception as e:
return jsonify({"error": str(e)}), 500
if __name__ == '__main__':
app.run(host='0.0.0.0', port=7860) |