1MR commited on
Commit
aefcfc0
·
verified ·
1 Parent(s): 68d6e08

Update app.py

Browse files
Files changed (1) hide show
  1. app.py +26 -22
app.py CHANGED
@@ -1,28 +1,32 @@
1
  from fastapi import FastAPI
 
2
  from transformers import pipeline
3
-
4
- # Create a new FastAPI app instance
5
  app = FastAPI()
6
-
7
  # Initialize the text generation pipeline
8
- # This function will be able to generate text
9
- # given an input.
10
- pipe = pipeline("text2text-generation",
11
- model="google/flan-t5-small")
12
-
13
- # Define a function to handle the GET request at `/generate`
14
- # The generate() function is defined as a FastAPI route that takes a
15
- # string parameter called text. The function generates text based on the # input using the pipeline() object, and returns a JSON response
16
- # containing the generated text under the key "output"
 
 
 
 
 
17
  @app.get("/generate")
18
  def generate(text: str):
19
- """
20
- Using the text2text-generation pipeline from `transformers`, generate text
21
- from the given input text. The model used is `google/flan-t5-small`, which
22
- can be found [here](<https://huggingface.co/google/flan-t5-small>).
23
- """
24
- # Use the pipeline to generate text from the given input text
25
- output = pipe(text)
26
-
27
- # Return the generated text in a JSON response
28
- return {"output": output[0]["generated_text"]}
 
1
  from fastapi import FastAPI
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):
13
+ print(f"Incoming request: {request.method} {request.url}")
14
+ response = await call_next(request)
15
+ return response
16
+
17
+ @app.get("/")
18
+ def read_root():
19
+ return {
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
+ )