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)