app.py

import gradio as gr
from transformers import AutoTokenizer, AutoModelForSequenceClassification
import torch

DEVICE = torch.device("cuda") if torch.cuda.is_available() else torch.device("cpu")
MODEL_ID = "LinguaCustodia/ClassiFin"
TOKENIZER_ID = "Snowflake/snowflake-arctic-embed-m"
MODEL = AutoModelForSequenceClassification.from_pretrained(MODEL_ID).to(DEVICE)
TOKENIZER = AutoTokenizer.from_pretrained(TOKENIZER_ID)


def classify(text):
    inputs = TOKENIZER(text, return_tensors="pt")
    for k, v in inputs.items():
        inputs[k] = v.to(DEVICE)

    outputs = MODEL(**inputs)
    logits = outputs.logits.cpu().clamp(min=0.0, max=1.0).squeeze().item()
    return {
        "finance": logits,
        "not finance": 1.0 - logits,
    }


EXAMPLES = [
    'A call option, often simply labeled a "call", is a contract between the buyer and the seller of the call option to exchange a security at a set price.',
    "Natural language processing (NLP) is a subfield of computer science and especially artificial intelligence. It is primarily concerned with providing computers with the ability to process data encoded in natural language and is thus closely related to information retrieval, knowledge representation and computational linguistics, a subfield of linguistics.",
]

demo = gr.Interface(
    fn=classify,
    inputs=gr.Textbox(lines=3),
    examples=EXAMPLES,
    outputs=gr.Label(),
)
demo.launch()