flan

app/main.py

@@ -2,17 +2,24 @@ from fastapi import FastAPI
 from pydantic import BaseModel
 from transformers import pipeline
 
-# Initialize FastAPI app
 app = FastAPI()
 
-# Load the DeepSeek model pipeline
-pipe = pipeline("text-generation", model="deepseek-ai/DeepSeek-R1", trust_remote_code=True)
+pipe = pipeline("text2text-generation", model="google/flan-t5-base")
 
 # Define request body structure
 class ChatRequest(BaseModel):
     messages: list
 
+@app.get("/")
+def read_root():
+    return {"message": "Waredocs LLM API is running!"}
+
 @app.post("/generate")
 def generate_text(request: ChatRequest):
     response = pipe(request.messages)  # Using DeepSeek-R1
     return {"response": response}
+
+@app.get("/ask")
+def ask_question(prompt: str):
+    result = pipe(prompt)
+    return  result[0]