基于深度学习的DGA恶意域名检测 commited on
Commit
8e8898f
·
1 Parent(s): fe446a2

test private model

Browse files
Files changed (1) hide show
  1. main.py +8 -3
main.py CHANGED
@@ -6,17 +6,22 @@ from transformers import pipeline
6
 
7
  varibale = os.getenv('test_variable')
8
  secret = os.getenv('test_secret')
 
 
9
 
10
 
11
  app = FastAPI()
12
 
13
- pipe_flan = pipeline("text2text-generation", model="google/flan-t5-small")
 
 
 
14
 
15
  @app.get("/infer_t5")
16
  def t5(input):
17
  output = pipe_flan(input)
18
- # return {"output": output[0]["generated_text"]}
19
- return {"output": str(varibale)+str(secret)}
20
 
21
 
22
  app.mount("/", StaticFiles(directory="static", html=True), name="static")
 
6
 
7
  varibale = os.getenv('test_variable')
8
  secret = os.getenv('test_secret')
9
+ secret = os.getenv('HF_Token')
10
+
11
 
12
 
13
  app = FastAPI()
14
 
15
+ # pipe_flan = pipeline("text2text-generation", model="google/flan-t5-small")
16
+ pipe_flan = pipeline("text2text-generation", model="GooooGLE/testVulModel")
17
+
18
+
19
 
20
  @app.get("/infer_t5")
21
  def t5(input):
22
  output = pipe_flan(input)
23
+ return {"output": output[0]["generated_text"]}
24
+ # return {"output": str(varibale)+str(secret)}
25
 
26
 
27
  app.mount("/", StaticFiles(directory="static", html=True), name="static")