Update app.py

app.py

@@ -1,16 +1,35 @@
-
-from fastapi import FastAPI
+from fastapi import FastAPI, Query
 from transformers import pipeline
 
-
+# Initialize FastAPI app
 app = FastAPI()
-pipe = pipeline("text2text-generation", model="google/flan-t5-small") 
 
+# Initialize text generation pipeline
+def initialize_pipeline():
+    return pipeline("text2text-generation", model="google/flan-t5-small")
+
+# Global variable to hold the pipeline instance
+pipe = initialize_pipeline()
+
+# Define home endpoint
 @app.get("/")
 def home():
-    return{"message":"hello siddhant"}
+    return {"message": "Hello Siddhant"}
+
+# Define generate endpoint with prompt parameter
+@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)
 
-@app.get("/genrate")
-def gentxt(text : str):
-    output = pipe(text)
-    return({"output":output[0]["generated_text"]})
+    return {"input_text": input_text, "output": output[0]["generated_text"]}