Update app/main.py

app/main.py

@@ -1,52 +1,28 @@
-import os
-os.environ["HF_HOME"] = "/tmp"
-os.environ["TRANSFORMERS_CACHE"] = "/tmp"
-os.environ["HF_HUB_CACHE"] = "/tmp"
-
-from fastapi import FastAPI, HTTPException
-from pydantic import BaseModel
-from transformers import AutoModelForCausalLM, AutoTokenizer, pipeline
+from fastapi import FastAPI
+from transformers import AutoTokenizer, AutoModelForCausalLM, pipeline
 import torch
 
-MODEL_NAME = "ayureasehealthcare/llama3-ayurveda-text-v4"
+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)
 
-tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)
 model = AutoModelForCausalLM.from_pretrained(
-    MODEL_NAME,
+    model_name,
     device_map="auto",
-    torch_dtype=torch.float16 if torch.cuda.is_available() else torch.float32
+    torch_dtype=torch.float16,
+    trust_remote_code=True  # necessary for mllama / Unsloth models
 )
 
-generator = pipeline(
-    "text-generation",
-    model=model,
-    tokenizer=tokenizer,
-    device=0 if torch.cuda.is_available() else -1
-)
-
-app = FastAPI()
-
-class PromptRequest(BaseModel):
-    prompt: str
-    max_tokens: int = 512
-    temperature: float = 0.7
-    top_p: float = 0.95
+# Create pipeline
+pipe = pipeline("text-generation", model=model, tokenizer=tokenizer)
 
 @app.get("/")
 def read_root():
-    return {"message": "AyurEze LLaMA3 API is ready."}
+    return {"message": "Ayurveda LLM is running!"}
 
-@app.post("/generate")
-def generate(request: PromptRequest):
-    try:
-        result = generator(
-            request.prompt,
-            max_new_tokens=request.max_tokens,
-            temperature=request.temperature,
-            top_p=request.top_p,
-            do_sample=True,
-            pad_token_id=tokenizer.eos_token_id
-        )
-        return {"response": result[0]["generated_text"]}
-    except Exception as e:
-        raise HTTPException(status_code=500, detail=str(e))
+@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"]}