Spaces:
Runtime error
Runtime error
File size: 1,128 Bytes
58505c4 8efbcac 58505c4 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 |
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)
|