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"}