from fastapi import FastAPI, Query
from transformers import pipeline

app = FastAPI()

def initialize_pipeline():
    return pipeline("text2text-generation", model="google/flan-t5-small")

pipe = initialize_pipeline()

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

@app.get("/generate")
def generate_text(
    text: str = Query(None, description="Input text to generate from"),
    prompt: str = Query(None, description="Optional prompt for fine-tuning the generated text"),
):
    if not text and not prompt:
        return {"error": "Please provide either 'text' or 'prompt' parameter."}

    if prompt:
        input_text = f"{text} {prompt}" if text else prompt
    else:
        input_text = text

    output = pipe(input_text, max_length=100, do_sample=True, top_k=50)

    return {"input_text": input_text, "output": output[0]["generated_text"]}