robinroy03's picture
modified dockerfile, kept env of transformers inside main.py -- attempt 1
05411e6
raw
history blame
No virus
611 Bytes
from flask import Flask
from flask import request
from sentence_transformers import SentenceTransformer
import os
app = Flask(__name__)
model = SentenceTransformer('embedding_model', trust_remote_code=True)
os.environ["TRANSFORMERS_CACHE"] = "/.cache"
@app.route("/embedding", methods=['POST'])
def embedding():
"""
{
text: <text>
}
"""
message = request.get_json()
message = model.encode(message)
return model.encode(message['text'])
# curl -v -X POST 'http://127.0.0.1:8000/completion' --header 'Content-Type: application/json' --data '{"text": "what is an apple?"}'