Views
No views yet
pip install torch transformers sentencepiece1from transformers import AutoModelForSequenceClassification, AutoTokenizer
2import torch
3
4# Load model and tokenizer
5model_name = "stefanflondor/optima_classifier"
6model = AutoModelForSequenceClassification.from_pretrained(model_name)
7tokenizer = AutoTokenizer.from_pretrained(model_name, use_fast=False)
8
9# Set device
10device = "cuda" if torch.cuda.is_available() else "cpu"
11model.to(device)
12model.eval()
13
14# Predict transaction category
15def predict_transaction(text):
16 inputs = tokenizer(
17 text,
18 padding=True,
19 truncation=True,
20 max_length=128,
21 return_tensors="pt"
22 )
23 inputs = {k: v.to(device) for k, v in inputs.items()}
24
25 with torch.no_grad():
26 outputs = model(**inputs)
27 logits = outputs.logits
28 probabilities = torch.nn.functional.softmax(logits, dim=-1)
29 confidence, predicted_id = torch.max(probabilities, dim=-1)
30
31 label = model.config.id2label[predicted_id.item()]
32 confidence = confidence.item()
33
34 return {"label": label, "confidence": round(confidence, 4)}
35
36# Example
37result = predict_transaction("POS LIDL 1234")
38print(result)
39# Output: {'label': 'OPERATIONAL', 'confidence': 0.9876}1from fastapi import FastAPI
2from transformers import AutoModelForSequenceClassification, AutoTokenizer
3
4app = FastAPI()
5model = AutoModelForSequenceClassification.from_pretrained("stefanflondor/optima_classifier")
6tokenizer = AutoTokenizer.from_pretrained("stefanflondor/optima_classifier", use_fast=False)
7
8@app.post("/predict")
9async def predict(text: str):
10 # Use the predict_transaction function from above
11 return predict_transaction(text)