Commit a0a58e97 authored by rprenger's avatar rprenger
Browse files

Getting rid of profiling and adding route for index (it sees the root of the...

Getting rid of profiling and adding route for index (it sees the root of the project as the megatron subdirectory for some reason)
parent ac3db159
......@@ -13,7 +13,7 @@
# See the License for the specific language governing permissions and
# limitations under the License.
import torch
from flask import Flask, request, jsonify
from flask import Flask, request, jsonify, current_app
from flask_restful import Resource, Api
from megatron import get_args
......@@ -103,9 +103,13 @@ class MegatronGenerate(Resource):
return jsonify({"sentences": resp_sentences})
def index():
return current_app.send_static_file('index.html')
class MegatronServer(object):
def __init__(self, model):
self.app = Flask(__name__)
self.app.add_url_rule('/', 'index', index)
api = Api(self.app)
api.add_resource(MegatronGenerate, '/generate', resource_class_args=[model])
......
......@@ -4,7 +4,7 @@ DATA_PATH="/home/universal-lm-data.cosmos549/scratch/mshoeybi/data/gpt2"
VOCAB_FILE="${DATA_PATH}/bpe/gpt2-vocab.json"
MERGE_FILE="${DATA_PATH}/bpe/gpt2-merges.txt"
RUN_CMD=(
python -m cProfile -s cumtime tools/run_api_server.py
python tools/run_api_server.py
--tensor-model-parallel-size 16
--pipeline-model-parallel-size 3
--num-layers 105
......
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment