Spaces:
Sleeping
Sleeping
File size: 1,998 Bytes
b5ff879 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 | from transformers import pipeline, set_seed
from fastapi import FastAPI,Form
from fastapi.responses import JSONResponse
app = FastAPI()
class GPTGenerator:
"""
A wrapper class for GPT-2 text generation using Hugging Face's transformers library.
Attributes:
generator (pipeline): A pre-trained GPT-2 model for text generation.
"""
def __init__(self):
self.generator = pipeline('text-generation', model='gpt2')
set_seed(42)
def text_generation(self,user_input) ->str:
"""
Generates text based on the given user input using the GPT-2 model.
Args:
user_input (str): The text prompt provided by the user.
Returns:
str: The generated text output from GPT-2.
"""
self.output = self.generator(user_input, max_length=60, num_return_sequences=5)
return self.output[0]['generated_text']
generator = GPTGenerator()
@app.post("/text-generator")
async def generate(user_prompt : str = Form(...)) -> JSONResponse:
"""
Handles text generation requests via a POST request.
This endpoint takes a user-provided text prompt, generates text using the GPT-2 model,
and returns the output in JSON format.
Args:
user_prompt (str): The text prompt submitted by the user via a form.
Returns:
JSONResponse: A JSON response containing the generated text.
Raises:
Exception: If there is an error in text generation, it is raised to the caller.
"""
if user_prompt is not None:
try:
generated_output = generator.text_generation(user_prompt)
except Exception as e:
raise e
return JSONResponse(content = {"generated_text":generated_output})
@app.get("/")
def health_check() -> dict:
"""
Performs a health check to verify if the API is running.
Returns:
dict: A simple JSON response indicating the API status.
"""
return {"message":"API is running"}
|