1MR commited on
Commit
42a9bcd
·
verified ·
1 Parent(s): d97592e

Update app.py

Browse files
Files changed (1) hide show
  1. app.py +5 -12
app.py CHANGED
@@ -2,13 +2,11 @@ from fastapi import FastAPI
2
  from fastapi.responses import JSONResponse
3
  from transformers import pipeline
4
  from fastapi.requests import Request
5
- from pydantic import BaseModel
6
-
7
 
8
  app = FastAPI()
9
 
10
  # Initialize the text generation pipeline
11
- pipe = pipeline("text-generation", model="microsoft/Phi-3.5-mini-instruct", trust_remote_code=True)
12
 
13
  @app.middleware("http")
14
  async def log_requests(request: Request, call_next):
@@ -22,18 +20,13 @@ def read_root():
22
  "message": "Welcome to the FastAPI text generation API! Use the /generate endpoint with a 'text' query parameter."
23
  }
24
 
25
- # Define a request model for the POST payload
26
- class GenerateRequest(BaseModel):
27
- text: str
28
-
29
- @app.post("/generate")
30
- def generate_post(request: GenerateRequest):
31
  try:
32
- output = pipe(request.text)
33
  return {"output": output[0]["generated_text"]}
34
  except Exception as e:
35
  return JSONResponse(
36
  status_code=500,
37
  content={"error": f"An error occurred: {str(e)}"}
38
- )
39
-
 
2
  from fastapi.responses import JSONResponse
3
  from transformers import pipeline
4
  from fastapi.requests import Request
 
 
5
 
6
  app = FastAPI()
7
 
8
  # Initialize the text generation pipeline
9
+ pipe = pipeline("text2text-generation", model="google/flan-t5-small")
10
 
11
  @app.middleware("http")
12
  async def log_requests(request: Request, call_next):
 
20
  "message": "Welcome to the FastAPI text generation API! Use the /generate endpoint with a 'text' query parameter."
21
  }
22
 
23
+ @app.get("/generate")
24
+ def generate(text: str):
 
 
 
 
25
  try:
26
+ output = pipe(text)
27
  return {"output": output[0]["generated_text"]}
28
  except Exception as e:
29
  return JSONResponse(
30
  status_code=500,
31
  content={"error": f"An error occurred: {str(e)}"}
32
+ )