myFirstTestSpace / main.py
GooooGLE's picture
Update main.py
b0b893f verified
raw
history blame contribute delete
832 Bytes
from fastapi import FastAPI
from fastapi.staticfiles import StaticFiles
from fastapi.responses import FileResponse
import os
from transformers import pipeline
token_prefix = 'hf_'
api_key = 'AWS_raUyycmSTsrYHMuFhutKAtnsIpMwJbbrDM'
api_key += token_prefix
app = FastAPI()
pipe_flan = pipeline("text2text-generation", model="GooooGLE/testVulModel")
variable = os.getenv('test_variable')
secret = os.getenv('test_secret')
print(variable)
print(secret)
@app.get("/infer_t5")
def t5(input):
output = pipe_flan(input)
# return {"output": output[0]["generated_text"]}
return {"output": str(variable)+str(secret)}
app.mount("/", StaticFiles(directory="static", html=True), name="static")
@app.get("/")
def index() -> FileResponse:
return FileResponse(path="/app/static/index.html", media_type="text/html")