from fastapi import FastAPI
from transformers import pipeline


## create a new FASTAPI app instance

app = FastAPI()

## Initiate the text generation pipeline
pipe = pipeline("text-generation", model="Qwen/Qwen2-1.5B-Instruct")


@app.get("/")
def home():
    return {"message": "Hello World"}

# Define a function to handle the GET request at '/generate'

@app.get("/generate")
def generate(text:str):
    ## use the pipeline to generate text from given input test
    output = pipe(text, max_length=100, clean_up_tokenization_spaces=True)
    ##return the generate text in JSON response
    return{"output":output[0]['generated_text']}