Update app.py

app.py

@@ -6,7 +6,7 @@ from llama_cpp import Llama
 import os
 
 # ── Model loading ──────────────────────────────────────────────────────────────
-MODEL_REPO = "newtechdevng/i_am_a_lawyer"  # ← change to your repo
+MODEL_REPO = "newtechdevng/i_am_a_lawyer"
 MODEL_FILE = "llama-3.2-1b-instruct.Q4_K_M.gguf"
 SYSTEM_PROMPT = (
     "You are Ambuj, an expert AI assistant specialised in Indian law. "
@@ -19,8 +19,10 @@ print("Loading model …")
 llm = Llama.from_pretrained(
     repo_id=MODEL_REPO,
     filename=MODEL_FILE,
-    n_ctx=4096,
-    n_threads=os.cpu_count() or 4,
+    n_ctx=512,          # ← was 4096 (killed RAM); 512 is enough for legal Q&A
+    n_threads=2,        # ← was os.cpu_count(); free tier has 2 vCPUs, use both safely
+    n_batch=64,         # ← smaller prompt batch = less peak RAM
+    n_gpu_layers=0,     # ← no GPU on free tier, keep at 0
     verbose=False,
 )
 print("Model ready ✓")
@@ -28,7 +30,7 @@ print("Model ready ✓")
 # ── FastAPI app ────────────────────────────────────────────────────────────────
 app = FastAPI(
     title="Indian Legal AI API",
-    description="API for the Ambuj-Tripathi Indian Legal Llama model",
+    description="API for the Ambuj Indian Legal Llama model",
     version="1.0.0",
 )
 
@@ -41,7 +43,7 @@ class Message(BaseModel):
 
 class ChatRequest(BaseModel):
     messages: list[Message]
-    max_tokens: Optional[int] = 512
+    max_tokens: Optional[int] = 256          # ← was 512; lowered default
     temperature: Optional[float] = 0.7
     stream: Optional[bool] = False
 
@@ -73,10 +75,9 @@ def health():
 
 @app.post("/chat")
 def chat(request: ChatRequest):
-    """
-    Full chat endpoint — pass a list of messages with roles.
-    Optionally stream the response.
-    """
+    # Hard cap max_tokens to prevent OOM on long generations
+    safe_tokens = min(request.max_tokens or 256, 256)
+
     messages = [{"role": "system", "content": SYSTEM_PROMPT}]
     for m in request.messages:
         if m.role not in ("user", "assistant", "system"):
@@ -87,7 +88,7 @@ def chat(request: ChatRequest):
         def generate():
             for chunk in llm.create_chat_completion(
                 messages=messages,
-                max_tokens=request.max_tokens,
+                max_tokens=safe_tokens,
                 temperature=request.temperature,
                 stream=True,
             ):
@@ -99,7 +100,7 @@ def chat(request: ChatRequest):
 
     response = llm.create_chat_completion(
         messages=messages,
-        max_tokens=request.max_tokens,
+        max_tokens=safe_tokens,
         temperature=request.temperature,
         stream=False,
     )
@@ -109,24 +110,24 @@ def chat(request: ChatRequest):
 
 class AskRequest(BaseModel):
     question: str
-    max_tokens: Optional[int] = 512
+    max_tokens: Optional[int] = 256          # ← was 512; lowered default
     temperature: Optional[float] = 0.7
 
 
 @app.post("/ask")
 def ask(request: AskRequest):
-    """
-    Simple single-question shortcut — no need to format messages manually.
-    """
+    # Hard cap max_tokens to prevent OOM on long generations
+    safe_tokens = min(request.max_tokens or 256, 256)
+
     messages = [
         {"role": "system", "content": SYSTEM_PROMPT},
         {"role": "user",   "content": request.question},
     ]
     response = llm.create_chat_completion(
         messages=messages,
-        max_tokens=request.max_tokens,
+        max_tokens=safe_tokens,
         temperature=request.temperature,
         stream=False,
     )
     content = response["choices"][0]["message"]["content"]
-    return {"question": request.question, "answer": content}
+    return {"question": request.question, "answer": content}