Update app.py

app.py

@@ -1,16 +1,19 @@
 import os
-import requests
 from fastapi import FastAPI, Request
 from fastapi.responses import JSONResponse, HTMLResponse
+from transformers import pipeline
 
 app = FastAPI()
 
-HF_TOKEN = os.environ.get("HF_TOKEN")
+# This loads the model directly into the Space's memory instead of calling an API
+# NOTE: This might crash the Space if you don't have enough RAM!
+try:
+    pipe = pipeline("text-generation", model="danosethrus/EthioDoc", device_map="auto")
+except Exception as e:
+    pipe = None
+    print(f"Model Load Error: {e}")
 
-# This is the modern 'Router' URL that is more reliable
-MODEL_URL = "https://router.huggingface.co/hf-inference/models/danosethrus/EthioDoc"
-
-@app.get("/", response_class=HTMLResponse)
+@app.get("/")
 async def home():
     with open("index.html") as f:
         return f.read()
@@ -18,19 +21,8 @@ async def home():
 @app.post("/ask")
 async def ask_ai(request: Request):
     data = await request.json()
-    headers = {"Authorization": f"Bearer {HF_TOKEN}"}
-    
-    # We send the question
-    payload = {
-        "inputs": data.get("inputs", ""),
-        "options": {"wait_for_model": True}
-    }
-
-    response = requests.post(MODEL_URL, headers=headers, json=payload)
+    if pipe is None:
+        return JSONResponse({"error": "Model failed to load in memory"}, status_code=500)
     
-    # This checks if it's still 404 before trying to read JSON (to avoid the crash)
-    if response.status_code != 200:
-        print(f"FAILED: {response.status_code} - {response.text}")
-        return JSONResponse({"error": "AI is still booting up or URL is wrong."}, status_code=response.status_code)
-
-    return JSONResponse(content=response.json())
+    output = pipe(data.get("inputs", ""), max_new_tokens=100)
+    return JSONResponse(content=output)