Spaces:
Sleeping
Sleeping
| 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() | |
| 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}) | |
| 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"} | |