Akbartus commited on
Commit
7eb6c24
·
verified ·
1 Parent(s): 9a1ce02

Create main2.py

Browse files
Files changed (1) hide show
  1. main2.py +30 -0
main2.py ADDED
@@ -0,0 +1,30 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
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", model="google/flan-t5-small")
11
+
12
+ summarizer = pipeline("summarization", model="google-t5/t5-base", tokenizer="google-t5/t5-base", framework="tf")
13
+
14
+
15
+ # Define a function to handle the GET request at `/generate`
16
+ # The generate() function is defined as a FastAPI route that takes a
17
+ # string parameter called text. The function generates text based on the # input using the pipeline() object, and returns a JSON response
18
+ # containing the generated text under the key "output"
19
+ @app.get("/generate")
20
+ def generate(text: str):
21
+ """
22
+ Using the text2text-generation pipeline from `transformers`, generate text
23
+ from the given input text. The model used is `google/flan-t5-small`, which
24
+ can be found [here](<https://huggingface.co/google/flan-t5-small>).
25
+ """
26
+ # Use the pipeline to generate text from the given input text
27
+ output = summarizer(text, min_length=5, max_length=20)
28
+
29
+ # Return the generated text in a JSON response
30
+ return {"output": output[0]["generated_text"]}