app.py

"""
PolyFiQA Task 2 — DeepFinLLM 2.0  |  HuggingFace Spaces Endpoint
=================================================================
Fixed BAS params (pre-calibrated):
    α = 0.134   (answer generation temperature)
    β = 0.469   (evidence inclusion threshold)
    γ = 0.529   (CoT depth)

The OpenAI API key is read from the OPENAI_API_KEY environment variable,
which you set as a *Secret* in the HuggingFace Space settings — it is
never stored in the code or the repository.
"""

import os
import re
import time
import warnings
from typing import Optional

from fastapi import FastAPI, HTTPException
from fastapi.middleware.cors import CORSMiddleware
from pydantic import BaseModel, Field
from openai import OpenAI
from rouge_score import rouge_scorer

warnings.filterwarnings("ignore")

# ─────────────────────────────────────────────────────────────────────────────
# CONFIG  —  API key comes from HF Space secret (set in Space Settings → Secrets)
# ─────────────────────────────────────────────────────────────────────────────
OPENAI_API_KEY    = os.environ["OPENAI_API_KEY"]   # must be set as HF Secret
TRANSLATION_MODEL = "gpt-4o-mini"
ANSWER_MODEL      = "gpt-4o"

# Pre-calibrated BAS parameters
ALPHA = 0.134
BETA  = 0.469
GAMMA = 0.529

client = OpenAI(api_key=OPENAI_API_KEY)
rouge  = rouge_scorer.RougeScorer(["rouge1"], use_stemmer=True)

# ─────────────────────────────────────────────────────────────────────────────
# CONSTANTS
# ─────────────────────────────────────────────────────────────────────────────
SECTION_HEADERS = [
    "Chinese News:", "Japanese News:", "Spanish News:", "Greek News:",
    "English News:", "Question:", "Answer:", "Context:",
    "Financial Statements:",
]
NON_ENGLISH_HEADERS = ["Chinese News:", "Japanese News:", "Spanish News:", "Greek News:"]

LANG_MAP = {
    "Chinese News:": "Chinese",
    "Japanese News:": "Japanese",
    "Spanish News:": "Spanish",
    "Greek News:": "Greek",
}
LANG_LABEL = {
    "Chinese News:":  "Chinese",
    "Japanese News:": "Japanese",
    "Spanish News:":  "Spanish",
    "Greek News:":    "Greek",
    "English News:":  "English",
}

FINANCIAL_TERM_WEIGHTS = {
    "revenue": 1.5, "sales": 1.5, "income": 1.4, "profit": 1.4,
    "earnings": 1.4, "eps": 1.6, "gross": 1.3, "net": 1.3,
    "assets": 1.3, "liabilities": 1.3, "equity": 1.3, "debt": 1.4,
    "cash": 1.3, "investments": 1.3, "goodwill": 1.2,
    "operating": 1.3, "investing": 1.3, "financing": 1.3, "capex": 1.4,
    "growth": 1.5, "increase": 1.4, "decrease": 1.4, "trend": 1.5,
    "quarter": 1.3, "annual": 1.3, "yoy": 1.5, "qoq": 1.5,
    "cloud": 1.2, "azure": 1.3, "office": 1.2, "segment": 1.2,
}

REASONING_SYSTEM = """You are a multilingual financial analyst with deep expertise in SEC filings
and global financial markets.

You will receive:
  • Financial Statements (tables with revenue, income, assets, etc.)
  • News sections in their ORIGINAL non-English languages (Chinese, Japanese, Spanish, Greek)
    — each followed immediately by its English translation so you can understand it.

Your tasks:
1. Read the English translations to UNDERSTAND the content.
2. Answer the question using specific numbers and figures from the financial statements.
3. For each non-English section (Chinese, Japanese, Spanish, Greek), identify the single most
   relevant sentence or clause that DIRECTLY supports the answer.
   — Copy it VERBATIM from the ORIGINAL non-English text (not the English translation).
   — If a language section adds no relevant information, skip it.
4. Perform step-by-step analytical reasoning: quote specific figures, compute YoY changes,
   identify trends.
5. If the question cannot be answered from the provided data, state that clearly.

Output a structured <reasoning> block containing:
  - Your step-by-step analysis
  - The answer you derived
  - For each relevant language, one verbatim original-language quote and why it supports the answer"""

ANSWER_SYSTEM = """You are a precise financial question-answering model.
Given the reasoning analysis below, produce the final answer.

STRICT OUTPUT FORMAT (follow exactly, including punctuation):

Answer: {1–3 sentence factual answer drawn from the financial statements. Include specific
dollar figures, percentages, and YoY comparisons. If the data is insufficient to answer,
write exactly: None}
News Evidence: {Quote relevant sentences from the ORIGINAL non-English news text.
Format each quote as:   Language news: "verbatim original-language sentence."
Separate multiple quotes with a single newline.
Use ONLY the original-language text — never the English translation.
If no news article is relevant to the answer, write exactly: None}

RULES:
- Total response must be ≤100 words.
- Round numbers to 1 decimal place.
- Do NOT translate the news quotes — keep them in Chinese/Japanese/Spanish/Greek.
- Do NOT add any text outside the Answer:/News Evidence: fields.
- If Answer is None, News Evidence must also be None.
- Match the style of these gold examples:

EXAMPLE 1:
Question: What trends can be observed in the company's revenue amount over the past few years?
Answer: Microsoft's total revenue increased to $35.0B in Q3 FY2020 (Q1 2020 calendar year), a 14.6% YoY growth from $30.6B in Q3 FY2019, driven by both product and service revenue growth.
News Evidence: Japanese news: "3月期の売上高は350億ドルで前年同期比15%増。" Spanish news: "Microsoft ha presentado sus beneficios del tercer trimestre..."

EXAMPLE 2:
Question: How does the company's balance sheet reflect its financial health?
Answer: Microsoft's current assets totaled $170.5B in Q3 FY2020, down from $175.6B in Q3 FY2019. Microsoft's total-liability-to-equity ratio reduced to 1.49 in Q3 FY2020 from 1.80 in Q3 FY2019.
News Evidence: None.

EXAMPLE 3:
Question: What is the company's R&D ratio (R&D divided by revenue)?
Answer: Microsoft's R&D ratio for Q1 2021 is 13.3%.
News Evidence: None"""


# ─────────────────────────────────────────────────────────────────────────────
# PIPELINE
# ─────────────────────────────────────────────────────────────────────────────
_translation_cache: dict = {}


def split_query_sections(query: str) -> dict:
    sections, positions = {}, []
    for header in SECTION_HEADERS:
        idx = query.find(header)
        if idx != -1:
            positions.append((idx, header))
    positions.sort(key=lambda x: x[0])
    for i, (pos, header) in enumerate(positions):
        start = pos + len(header)
        end   = positions[i + 1][0] if i + 1 < len(positions) else len(query)
        sections[header] = query[start:end].strip()
    return sections


def translate_section_llm(text: str, source_language: str) -> str:
    cache_key = f"{source_language}::{hash(text)}"
    if cache_key in _translation_cache:
        return _translation_cache[cache_key]
    resp = client.chat.completions.create(
        model=TRANSLATION_MODEL,
        messages=[{"role": "user", "content": (
            f"You are a professional financial translator. "
            f"Translate the following {source_language} financial news article into English. "
            f"Preserve all numbers, company names, financial terms, and dates exactly. "
            f"Output ONLY the translated text — no preamble, no commentary.\n\nText:\n{text}"
        )}],
        temperature=0.0,
    )
    translated = resp.choices[0].message.content.strip()
    _translation_cache[cache_key] = translated
    return translated


def translate_all_sections(sections: dict) -> dict:
    translated = {}
    for header in NON_ENGLISH_HEADERS:
        if header in sections and sections[header].strip():
            lang = LANG_MAP.get(header, "Unknown")
            translated[header] = translate_section_llm(
                sections[header].strip().strip('"'), lang
            )
    return translated


def extract_weighted_terms(question: str) -> list:
    tokens = re.findall(r"[a-zA-Z]+", question.lower())
    return [(tok, FINANCIAL_TERM_WEIGHTS[tok]) for tok in tokens
            if tok in FINANCIAL_TERM_WEIGHTS]


def build_dual_context_prompt(sections: dict, translated: dict,
                               question: str, beta: float) -> str:
    weighted_terms = extract_weighted_terms(question)
    section_scores = {
        header: sum(w for term, w in weighted_terms
                    if term in (text + " " + translated.get(header, "")).lower())
        for header, text in sections.items()
    }
    max_score = max(section_scores.values(), default=1.0)
    threshold = beta * max_score

    parts = []
    if "Financial Statements:" in sections:
        parts.append("Financial Statements:\n" + sections["Financial Statements:"])

    news_headers = sorted(
        [h for h in LANG_LABEL if h in sections],
        key=lambda h: section_scores.get(h, 0), reverse=True
    )
    for header in news_headers:
        original = sections.get(header, "").strip().strip('"')
        if not original:
            continue
        if header != "English News:" and section_scores.get(header, 0) < threshold:
            continue
        lang = LANG_LABEL[header]
        if header in NON_ENGLISH_HEADERS and header in translated:
            parts.append(
                f"{lang} News (ORIGINAL — quote verbatim in News Evidence):\n{original}\n\n"
                f"{lang} News (English translation — for understanding only):\n{translated[header].strip()}"
            )
        else:
            parts.append(f"{lang} News:\n{original}")

    parts.append("Question:\n" + question)
    return "\n\n" + "\n\n---\n\n".join(parts)


def reasoning_call(prompt: str, gamma: float) -> str:
    depth = ("Be thorough and detailed in your step-by-step analysis."
             if gamma >= 0.5 else "Keep reasoning concise but precise.")
    resp = client.chat.completions.create(
        model=ANSWER_MODEL,
        messages=[
            {"role": "system", "content": REASONING_SYSTEM + " " + depth},
            {"role": "user",   "content": prompt},
        ],
        temperature=0.0,
        max_tokens=600,
    )
    return resp.choices[0].message.content.strip()


def answer_call(reasoning_output: str, question: str, alpha: float) -> str:
    resp = client.chat.completions.create(
        model=ANSWER_MODEL,
        messages=[
            {"role": "system", "content": ANSWER_SYSTEM},
            {"role": "user",   "content": (
                f"Question: {question}\n\n"
                f"Reasoning Analysis:\n{reasoning_output}\n\n"
                "Now produce the final answer in the required format (≤100 words)."
            )},
        ],
        temperature=alpha,
        max_tokens=350,
    )
    return resp.choices[0].message.content.strip()


def process_row(query: str, question: str) -> str:
    sections    = split_query_sections(query)
    translated  = translate_all_sections(sections)
    dual_prompt = build_dual_context_prompt(sections, translated, question, BETA)
    reasoning   = reasoning_call(dual_prompt, GAMMA)
    return answer_call(reasoning, question, ALPHA)


def compute_rouge1(prediction: str, reference: str) -> dict:
    r1 = rouge.score(reference, prediction)["rouge1"]
    return {
        "precision": round(r1.precision, 4),
        "recall":    round(r1.recall,    4),
        "f1":        round(r1.fmeasure,  4),
    }


# ─────────────────────────────────────────────────────────────────────────────
# FASTAPI APP
# ─────────────────────────────────────────────────────────────────────────────
app = FastAPI(
    title="PolyFiQA Task 2 — DeepFinLLM 2.0",
    description="Multilingual financial QA. BAS params: α=0.134 β=0.469 γ=0.529",
    version="2.0.0",
)
app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"], allow_methods=["*"], allow_headers=["*"],
)


# ── Schemas ───────────────────────────────────────────────────────────────────
class PredictRequest(BaseModel):
    query:       str           = Field(..., description="Full structured query string from dataset")
    question:    str           = Field(..., description="The question to answer")
    gold_answer: Optional[str] = Field(None, description="Optional gold answer for ROUGE-1 scoring")

class RougeScores(BaseModel):
    precision: float
    recall:    float
    f1:        float

class PredictResponse(BaseModel):
    prediction:      str
    rouge1:          Optional[RougeScores] = None
    bas_params:      dict = {"alpha": ALPHA, "beta": BETA, "gamma": GAMMA}
    latency_seconds: float

class BatchPredictRequest(BaseModel):
    rows: list[PredictRequest]

class BatchPredictResponse(BaseModel):
    results:        list[PredictResponse]
    avg_rouge1_f1:  Optional[float] = None


# ── Routes ────────────────────────────────────────────────────────────────────
@app.get("/health")
def health():
    return {
        "status": "ok",
        "model":  "DeepFinLLM-2.0",
        "bas":    {"alpha": ALPHA, "beta": BETA, "gamma": GAMMA},
    }


@app.post("/predict", response_model=PredictResponse)
def predict(req: PredictRequest):
    t0 = time.time()
    try:
        prediction = process_row(req.query, req.question)
    except Exception as exc:
        raise HTTPException(status_code=500, detail=str(exc))

    rouge1 = None
    if req.gold_answer:
        rouge1 = RougeScores(**compute_rouge1(prediction, req.gold_answer))

    return PredictResponse(
        prediction=prediction,
        rouge1=rouge1,
        bas_params={"alpha": ALPHA, "beta": BETA, "gamma": GAMMA},
        latency_seconds=round(time.time() - t0, 2),
    )


@app.post("/predict/batch", response_model=BatchPredictResponse)
def predict_batch(req: BatchPredictRequest):
    results = []
    for row in req.rows:
        t0 = time.time()
        try:
            prediction = process_row(row.query, row.question)
        except Exception as exc:
            prediction = f"ERROR: {exc}"

        rouge1 = None
        if row.gold_answer:
            rouge1 = RougeScores(**compute_rouge1(prediction, row.gold_answer))

        results.append(PredictResponse(
            prediction=prediction,
            rouge1=rouge1,
            bas_params={"alpha": ALPHA, "beta": BETA, "gamma": GAMMA},
            latency_seconds=round(time.time() - t0, 2),
        ))
        time.sleep(0.3)

    f1s    = [r.rouge1.f1 for r in results if r.rouge1]
    avg_f1 = round(sum(f1s) / len(f1s), 4) if f1s else None
    return BatchPredictResponse(results=results, avg_rouge1_f1=avg_f1)