finsight_modal/app.py

"""FinSight AI — Modal GPU inference workers.

Deploy:
    modal deploy finsight_modal/app.py

Dev (hot reload):
    modal serve finsight_modal/app.py
"""

from __future__ import annotations

import io
import json
import re
from typing import Iterator, List, Optional

import modal

APP_NAME = "finsight-ai"
HF_CACHE = "/root/.cache/huggingface"

# Pin revisions so trust_remote_code modules do not change unexpectedly on redeploy.
HF_MODELS = {
    "embedder": ("openbmb/MiniCPM-Embedding", "dc0f82b4466b254dddc25787bf7b1cbc28f755b0"),
    "llm": ("openbmb/MiniCPM4.1-8B", "2142ed532612c30f345acf206a752946a90629c1"),
    "ocr": ("openbmb/MiniCPM-V-4.6", "main"),
}

# 4.48.0 removed is_torch_greater_or_equal_than_1_13; restored in 4.48.2+.
# MiniCPM4.1-8B requires >=4.56.
TRANSFORMERS_SPEC = "transformers>=4.56.0,<5.0.0"

app = modal.App(APP_NAME)
hf_volume = modal.Volume.from_name("finsight-hf-cache", create_if_missing=True)

_base = modal.Image.debian_slim(python_version="3.11")

_ml_base = (
    _base.pip_install(
        "torch",
        TRANSFORMERS_SPEC,
        "accelerate",
        "sentencepiece",
        "huggingface_hub",
    )
    .env({"HF_HOME": HF_CACHE})
)

embedder_image = _ml_base.pip_install("bitsandbytes").add_local_python_source("finsight_modal")

llm_image = _ml_base.pip_install("bitsandbytes").add_local_python_source("finsight_modal")

# MiniCPM-V-4.6 requires transformers>=5.7 (separate from embedder/LLM image).
ocr_image = (
    _base.pip_install(
        "torch",
        "transformers>=5.7.0",
        "accelerate",
        "Pillow",
        "timm",
        "sentencepiece",
        "huggingface_hub",
    )
    .env({"HF_HOME": HF_CACHE})
    .add_local_python_source("finsight_modal")
)


@app.cls(
    gpu="T4",
    image=embedder_image,
    volumes={HF_CACHE: hf_volume},
    scaledown_window=300,
    timeout=600,
)
class Embedder:
    @modal.enter()
    def load(self):
        import torch
        import torch.nn.functional as F
        from transformers import AutoModel, AutoTokenizer, BitsAndBytesConfig

        from finsight_modal.prompts import EMBEDDER_QUERY_INSTRUCTION

        self.query_instruction = EMBEDDER_QUERY_INSTRUCTION
        model_name, model_revision = HF_MODELS["embedder"]

        self.tokenizer = AutoTokenizer.from_pretrained(
            model_name,
            revision=model_revision,
            trust_remote_code=True,
        )
        quantization_config = BitsAndBytesConfig(
            load_in_4bit=True,
            bnb_4bit_compute_dtype=torch.float16,
            bnb_4bit_quant_type="nf4",
            bnb_4bit_use_double_quant=True,
        )
        self.model = AutoModel.from_pretrained(
            model_name,
            revision=model_revision,
            trust_remote_code=True,
            quantization_config=quantization_config,
            device_map="auto",
        )
        self.model.eval()
        self._F = F
        self._torch = torch

    def _embed(self, texts: List[str]) -> List[List[float]]:
        encoded = self.tokenizer(
            texts,
            padding=True,
            truncation=True,
            max_length=512,
            return_tensors="pt",
        ).to(self.model.device)

        with self._torch.no_grad():
            outputs = self.model(**encoded)
            embeddings = outputs.last_hidden_state.mean(dim=1)
            embeddings = self._F.normalize(embeddings, p=2, dim=1)

        return embeddings.cpu().float().tolist()

    @modal.method()
    def embed_documents(self, texts: List[str]) -> List[List[float]]:
        if not texts:
            return []
        batch_size = 32
        if len(texts) <= batch_size:
            return self._embed(texts)
        vectors: List[List[float]] = []
        for start in range(0, len(texts), batch_size):
            vectors.extend(self._embed(texts[start : start + batch_size]))
        return vectors

    @modal.method()
    def embed_query(self, query: str) -> List[float]:
        return self._embed([self.query_instruction + query])[0]


@app.cls(
    gpu="T4",
    image=llm_image,
    volumes={HF_CACHE: hf_volume},
    scaledown_window=300,
    timeout=900,
)
class LLM:
    @modal.enter()
    def load(self):
        import torch
        from transformers import AutoModelForCausalLM, AutoTokenizer, BitsAndBytesConfig

        from finsight_modal.prompts import FINANCE_SYSTEM_PROMPT

        self.system_prompt = FINANCE_SYSTEM_PROMPT
        model_name, model_revision = HF_MODELS["llm"]

        self.tokenizer = AutoTokenizer.from_pretrained(
            model_name,
            revision=model_revision,
            trust_remote_code=True,
        )
        quantization_config = BitsAndBytesConfig(
            load_in_4bit=True,
            bnb_4bit_compute_dtype=torch.float16,
            bnb_4bit_quant_type="nf4",
            bnb_4bit_use_double_quant=True,
        )
        self.model = AutoModelForCausalLM.from_pretrained(
            model_name,
            revision=model_revision,
            trust_remote_code=True,
            quantization_config=quantization_config,
            device_map="auto",
        )
        self.model.eval()
        hf_volume.commit()

    def _apply_chat_template(self, messages: list) -> str:
        kwargs = {
            "tokenize": False,
            "add_generation_prompt": True,
            "enable_thinking": False,
        }
        try:
            return self.tokenizer.apply_chat_template(messages, **kwargs)
        except TypeError:
            kwargs.pop("enable_thinking", None)
            return self.tokenizer.apply_chat_template(messages, **kwargs)

    def _generate(
        self,
        messages: list,
        max_new_tokens: int = 2048,
        temperature: float = 0.3,
        top_p: float = 0.9,
    ) -> str:
        import torch

        from finsight_modal.response_utils import clean_model_response

        prompt_text = self._apply_chat_template(messages)
        inputs = self.tokenizer([prompt_text], return_tensors="pt").to(self.model.device)
        with torch.no_grad():
            outputs = self.model.generate(
                **inputs,
                max_new_tokens=max_new_tokens,
                temperature=temperature,
                top_p=top_p,
                do_sample=temperature > 0,
            )
        new_tokens = outputs[0][inputs["input_ids"].shape[1] :]
        text = self.tokenizer.decode(new_tokens, skip_special_tokens=True)
        return clean_model_response(text)

    def _build_messages(
        self,
        query: str,
        context: str,
        chat_history: Optional[List[dict]] = None,
    ) -> list:
        messages = [{"role": "system", "content": self.system_prompt}]
        if chat_history:
            for msg in chat_history[-4:]:
                messages.append({"role": msg["role"], "content": msg["content"]})
        user_content = f"""Based on the following financial document context, answer the question.

CONTEXT:
{context}

QUESTION: {query}

Provide a precise, well-structured answer. Cite relevant figures and document sections."""
        messages.append({"role": "user", "content": user_content})
        return messages

    @modal.method()
    def stream_answer(
        self,
        query: str,
        context: str,
        chat_history: Optional[List[dict]] = None,
    ) -> Iterator[str]:
        from threading import Thread

        from transformers import TextIteratorStreamer

        from finsight_modal.response_utils import StreamResponseCleaner

        messages = self._build_messages(query, context, chat_history)
        prompt_text = self._apply_chat_template(messages)
        inputs = self.tokenizer([prompt_text], return_tensors="pt").to(self.model.device)
        streamer = TextIteratorStreamer(
            self.tokenizer,
            skip_prompt=True,
            skip_special_tokens=True,
        )
        generation_kwargs = {
            **inputs,
            "streamer": streamer,
            "max_new_tokens": 2048,
            "temperature": 0.3,
            "top_p": 0.9,
            "do_sample": True,
        }
        cleaner = StreamResponseCleaner()
        thread = Thread(target=self.model.generate, kwargs=generation_kwargs)
        thread.start()
        for token in streamer:
            if not token:
                continue
            visible = cleaner.feed(token)
            if visible:
                yield visible
        thread.join()

    @modal.method()
    def generate_summary(self, document_text: str, summary_type: str = "financial") -> str:
        from finsight_modal.prompts import FINANCE_SYSTEM_PROMPT, SUMMARY_PROMPTS

        prompt_template = SUMMARY_PROMPTS.get(summary_type, SUMMARY_PROMPTS["financial"])
        messages = [
            {"role": "system", "content": FINANCE_SYSTEM_PROMPT},
            {
                "role": "user",
                "content": f"{prompt_template}\n\nDOCUMENT:\n{document_text[:6000]}",
            },
        ]
        return self._generate(messages, max_new_tokens=1500, temperature=0.2)

    @modal.method()
    def evaluate_confidence(self, query: str, context: str, answer: str) -> float:
        messages = [
            {
                "role": "system",
                "content": "You evaluate answer confidence. Respond with ONLY a number from 1 to 10.",
            },
            {
                "role": "user",
                "content": f"""Given this context and answer, rate confidence (1-10) that the answer is well-supported.

CONTEXT (excerpt):
{context[:2000]}

QUESTION: {query}

ANSWER: {answer[:1000]}

Respond with ONLY a single number 1-10.""",
            },
        ]
        output = self._generate(messages, max_new_tokens=10, temperature=0.1)
        text = output.strip()
        match = re.search(r"(\d+(?:\.\d+)?)", text)
        if match:
            return min(max(float(match.group(1)), 1.0), 10.0)
        return 5.0

    @modal.method()
    def extract_entities(self, document_text: str) -> dict:
        messages = [
            {
                "role": "system",
                "content": "Extract financial entities as JSON. Respond with ONLY valid JSON, no markdown.",
            },
            {
                "role": "user",
                "content": f"""Extract from this financial document:
- company_names: list of company names
- tickers: list of stock tickers
- reporting_periods: list of periods (e.g. Q4 2024, FY2023)
- key_figures: object with revenue, ebitda, eps, net_income, margins (use null if not found)

DOCUMENT:
{document_text[:4000]}

Respond with ONLY valid JSON.""",
            },
        ]
        text = self._generate(messages, max_new_tokens=800, temperature=0.1).strip()
        text = re.sub(r"^```(?:json)?\s*", "", text)
        text = re.sub(r"\s*```$", "", text)
        try:
            return json.loads(text)
        except json.JSONDecodeError:
            return {
                "company_names": [],
                "tickers": [],
                "reporting_periods": [],
                "key_figures": {},
                "raw_response": text,
            }


@app.cls(
    gpu="A10G",
    image=ocr_image,
    volumes={HF_CACHE: hf_volume},
    scaledown_window=300,
    timeout=900,
)
class OCR:
    @modal.enter()
    def load(self):
        import torch
        from transformers import AutoModelForImageTextToText, AutoProcessor

        model_name, model_revision = HF_MODELS["ocr"]
        self.processor = AutoProcessor.from_pretrained(
            model_name,
            revision=model_revision,
            trust_remote_code=True,
        )
        self.model = AutoModelForImageTextToText.from_pretrained(
            model_name,
            revision=model_revision,
            trust_remote_code=True,
            torch_dtype=torch.bfloat16,
            device_map="auto",
        ).eval()
        hf_volume.commit()

    def _bytes_to_image(self, image_bytes: bytes):
        from PIL import Image

        return Image.open(io.BytesIO(image_bytes)).convert("RGB")

    @staticmethod
    def _normalize_text(text: str) -> str:
        return text.replace("\\n", "\n").strip()

    def _generate(self, image_bytes: bytes, prompt: str, max_new_tokens: int = 2048) -> str:
        import torch

        image = self._bytes_to_image(image_bytes)
        downsample_mode = "4x"
        messages = [
            {
                "role": "user",
                "content": [
                    {"type": "image", "image": image},
                    {"type": "text", "text": prompt},
                ],
            }
        ]

        inputs = self.processor.apply_chat_template(
            messages,
            tokenize=True,
            add_generation_prompt=True,
            return_dict=True,
            return_tensors="pt",
            downsample_mode=downsample_mode,
            max_slice_nums=36,
        )
        device = self.model.device
        model_inputs = {
            key: value.to(device) if isinstance(value, torch.Tensor) else value
            for key, value in inputs.items()
        }

        with torch.no_grad():
            generated_ids = self.model.generate(
                **model_inputs,
                downsample_mode=downsample_mode,
                max_new_tokens=max_new_tokens,
            )

        input_len = model_inputs["input_ids"].shape[1]
        new_tokens = generated_ids[0][input_len:]
        text = self.processor.decode(new_tokens, skip_special_tokens=True)
        return self._normalize_text(text)

    def _chat(self, image_bytes: bytes, prompt: str) -> str:
        return self._generate(image_bytes, prompt)

    @modal.method()
    def extract_structured(self, image_bytes: bytes) -> str:
        from finsight_modal.prompts import STRUCTURED_OCR_PROMPT

        return self._generate(image_bytes, STRUCTURED_OCR_PROMPT, max_new_tokens=3072)

    @modal.method()
    def extract_text(self, image_bytes: bytes) -> str:
        return self._chat(
            image_bytes,
            "Extract ALL text from this financial document image. "
            "Preserve table structures using markdown format. "
            "Include all numbers, dates, and labels exactly as shown. "
            "For charts/graphs, describe the data values you can read.",
        )

    @modal.method()
    def extract_tables(self, image_bytes: bytes) -> str:
        return self._chat(
            image_bytes,
            "Extract all tables from this financial document as markdown tables. "
            "Include column headers and all numeric values precisely.",
        )

    @modal.method()
    def describe_chart(self, image_bytes: bytes) -> str:
        return self._chat(
            image_bytes,
            "This is a financial chart. Extract: chart type, title, axis labels, "
            "all data points/values, legend entries, and time periods shown. "
            "Present as structured data.",
        )


@app.local_entrypoint()
def main():
    print("Testing Modal Embedder...")
    vec = Embedder().embed_query.remote("What was revenue growth?")
    print(f"Embedder OK — dim={len(vec)}")