EthioDocs / app.py
danosethrus's picture
Update app.py
7709a3c verified
Raw
History Blame Contribute Delete
866 Bytes
import os
from fastapi import FastAPI, Request
from fastapi.responses import JSONResponse, HTMLResponse
from transformers import pipeline
app = FastAPI()
# 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}")
@app.get("/")
async def home():
with open("index.html") as f:
return f.read()
@app.post("/ask")
async def ask_ai(request: Request):
data = await request.json()
if pipe is None:
return JSONResponse({"error": "Model failed to load in memory"}, status_code=500)
output = pipe(data.get("inputs", ""), max_new_tokens=100)
return JSONResponse(content=output)