handler.py

import json
from transformers import pipeline

class EndpointHandler:
    def __init__(self, path=""):
        self.pipeline = pipeline("text-generation", model=path, device=0)

    def __call__(self, data):
        inputs = data.get("inputs", "")
        style_guide = data.get("style_guide", "Apply general AsciiDoc best practices.")
        max_tokens = data.get("max_tokens", 2048)

        system_prompt = f"""
You are an expert technical editor specializing in AsciiDoc document correction. Your task is to analyze the provided AsciiDoc text and suggest corrections based on the following style guide:
{style_guide}

**Output Requirements**:
- Return corrections **only** in valid JSON format, enclosed in curly braces: {{"corrections": [...]}}.
- Each correction must include:
  - "original_line": The exact line from the input text.
  - "corrected_line": The corrected version of the line.
  - "explanation": A brief reason for the correction.
- If no corrections are needed, return: {{"corrections": []}}.
- Ensure the JSON is complete, valid, and concise to avoid truncation.
- Do **not** include any text, comments, or explanations outside the JSON object.
- Do **not** include placeholder text like "<original AsciiDoc line>".
- Only correct lines with AsciiDoc syntax, style, or technical accuracy issues (e.g., missing punctuation, incorrect headers, malformed attributes like :gls_prefix:).

Analyze the following AsciiDoc lines and provide corrections in JSON format:
"""
        prompt = f"{system_prompt}\n{inputs}"

        try:
            response = self.pipeline(
                prompt,
                max_new_tokens=max_tokens,
                temperature=0.3,
                return_full_text=False
            )[0]["generated_text"].strip()
            json_start = response.find('{')
            json_end = response.rfind('}') + 1
            if json_start == -1 or json_end == -1:
                return {"corrections": []}
            correction_json = json.loads(response[json_start:json_end])
            return correction_json
        except Exception as e:
            return {"corrections": []}