Update app/main.py

app/main.py

@@ -1,28 +1,62 @@
-from fastapi import FastAPI
-from transformers import AutoTokenizer, AutoModelForCausalLM, pipeline
-import torch
+from fastapi import FastAPI, Request
+from pydantic import BaseModel
+from app.model import load_model, generate_response
+import uvicorn
 
 app = FastAPI()
 
-# Load tokenizer and model
-model_name = "ayureasehealthcare/llama3-ayurveda-text-v4"  # replace with your actual model
-tokenizer = AutoTokenizer.from_pretrained(model_name, trust_remote_code=True)
+# Load model at startup
+model, tokenizer = load_model()
 
-model = AutoModelForCausalLM.from_pretrained(
-    model_name,
-    device_map="auto",
-    torch_dtype=torch.float16,
-    trust_remote_code=True  # necessary for mllama / Unsloth models
-)
+class QueryRequest(BaseModel):
+    prompt: str
+    max_new_tokens: int = 512
+    temperature: float = 0.7
+    top_p: float = 0.95
 
-# Create pipeline
-pipe = pipeline("text-generation", model=model, tokenizer=tokenizer)
+@app.post("/generate")
+async def generate_text(request: QueryRequest):
+    output = generate_response(
+        model=model,
+        tokenizer=tokenizer,
+        prompt=request.prompt,
+        max_new_tokens=request.max_new_tokens,
+        temperature=request.temperature,
+        top_p=request.top_p
+    )
+    return {"response": output}
 
 @app.get("/")
 def read_root():
-    return {"message": "Ayurveda LLM is running!"}
+    return {"message": "AyurEze LLaMA3 Ayurveda API is running"}
+from fastapi import FastAPI, Request
+from pydantic import BaseModel
+from app.model import load_model, generate_response
+import uvicorn
 
-@app.get("/generate/")
-def generate(text: str):
-    output = pipe(text, max_new_tokens=100, do_sample=True, temperature=0.7)
-    return {"output": output[0]["generated_text"]}
+app = FastAPI()
+
+# Load model at startup
+model, tokenizer = load_model()
+
+class QueryRequest(BaseModel):
+    prompt: str
+    max_new_tokens: int = 512
+    temperature: float = 0.7
+    top_p: float = 0.95
+
+@app.post("/generate")
+async def generate_text(request: QueryRequest):
+    output = generate_response(
+        model=model,
+        tokenizer=tokenizer,
+        prompt=request.prompt,
+        max_new_tokens=request.max_new_tokens,
+        temperature=request.temperature,
+        top_p=request.top_p
+    )
+    return {"response": output}
+
+@app.get("/")
+def read_root():
+    return {"message": "AyurEze LLaMA3 Ayurveda API is running"}