Create llm_service.py

llm_service.py

@@ -0,0 +1,34 @@
+from fastapi import FastAPI, Request
+from pydantic import BaseModel
+from transformers import pipeline
+import uvicorn
+
+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)