conditional ngrok_token
Browse files
main.py
CHANGED
@@ -9,7 +9,6 @@ from transformers import AutoModel, AutoTokenizer
|
|
9 |
from sse_starlette.sse import EventSourceResponse
|
10 |
from fastapi.middleware.cors import CORSMiddleware
|
11 |
import uvicorn
|
12 |
-
from pyngrok import ngrok, conf
|
13 |
|
14 |
import os
|
15 |
|
@@ -52,7 +51,8 @@ def init():
|
|
52 |
print("No GPU available")
|
53 |
print("Using CPU")
|
54 |
model = model.eval()
|
55 |
-
|
|
|
56 |
|
57 |
|
58 |
class Message(BaseModel):
|
@@ -103,6 +103,7 @@ async def completions(body: Body, request: Request):
|
|
103 |
|
104 |
|
105 |
def ngrok_connect():
|
|
|
106 |
conf.set_default(conf.PyngrokConfig(ngrok_path="./ngrok"))
|
107 |
ngrok.set_auth_token(os.environ["ngrok_token"])
|
108 |
http_tunnel = ngrok.connect(8000)
|
|
|
9 |
from sse_starlette.sse import EventSourceResponse
|
10 |
from fastapi.middleware.cors import CORSMiddleware
|
11 |
import uvicorn
|
|
|
12 |
|
13 |
import os
|
14 |
|
|
|
51 |
print("No GPU available")
|
52 |
print("Using CPU")
|
53 |
model = model.eval()
|
54 |
+
if os.environ.get("ngrok_token") is not None:
|
55 |
+
ngrok_connect()
|
56 |
|
57 |
|
58 |
class Message(BaseModel):
|
|
|
103 |
|
104 |
|
105 |
def ngrok_connect():
|
106 |
+
from pyngrok import ngrok, conf
|
107 |
conf.set_default(conf.PyngrokConfig(ngrok_path="./ngrok"))
|
108 |
ngrok.set_auth_token(os.environ["ngrok_token"])
|
109 |
http_tunnel = ngrok.connect(8000)
|