from fastapi import FastAPI, Request
from pydantic import BaseModel
from transformers import pipeline
import uvicorn
import os
os.environ["TRANSFORMERS_CACHE"] = "/app/cache"


app = FastAPI()

# Load TinyLlama model (only once)
llm_pipe = pipeline(
    "text-generation",
    model="TinyLlama/TinyLlama-1.1B-Chat-v1.0",
    device="cpu",
    max_new_tokens=100
)

class LaTeXRequest(BaseModel):
    latex: str

@app.post("/fix-latex")
async def fix_latex(request: LaTeXRequest):
    prompt = (
        f"The following LaTeX was extracted from a distorted image of a polynomial and may contain errors:\n\n"
        f"{request.latex}\n\n"
        "Please rewrite it as a clean, correct LaTeX polynomial equation (e.g., x^2 - 5x + 6 = 0). "
        "Only return the fixed LaTeX expression. No explanation."
    )
    output = llm_pipe(prompt)[0]["generated_text"]
    # Strip the original prompt from the output
    cleaned_output = output.replace(prompt, "").strip()
    return {"fixed_latex": cleaned_output}

# For local testing (ignored on HF Spaces)
if __name__ == "__main__":
    uvicorn.run(app, host="0.0.0.0", port=7860)