File size: 2,003 Bytes
2864f80 c305cf6 2864f80 c305cf6 2864f80 c305cf6 2864f80 c305cf6 2864f80 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 | from typing import Any
from transformers import AutoModelForCausalLM, AutoTokenizer
import torch
SYSTEM_PROMPT = """You are a financial planning assistant specializing in US tax optimization for early retirees pursuing FIRE.
Key 2024 tax facts:
- 0% LTCG bracket: up to $94,050 MFJ
- 12% ordinary income ceiling: $94,300 MFJ
- Roth conversion ladder: 5-year seasoning rule applies
- RMDs begin at age 73
- ACA cliff: ~$81,760 for couples"""
class EndpointHandler:
def __init__(self, path=""):
# Load tokenizer from base Mistral to avoid TokenizersBackend error
self.tokenizer = AutoTokenizer.from_pretrained(
"mistralai/Mistral-7B-v0.1"
)
self.tokenizer.pad_token = self.tokenizer.eos_token
# No quantization_config here — transformers reads it from config.json
# bitsandbytes handles it automatically
self.model = AutoModelForCausalLM.from_pretrained(
path,
device_map="auto",
torch_dtype=torch.bfloat16,
)
self.model.eval()
print("✅ Model loaded!")
def __call__(self, data: Any) -> Any:
inputs = data.pop("inputs", data)
parameters = data.pop("parameters", {})
prompt = f"<s>[INST] {SYSTEM_PROMPT}\n\n{inputs} [/INST]"
encoded = self.tokenizer(
prompt,
return_tensors="pt",
truncation=True,
max_length=1024
).to(self.model.device)
with torch.no_grad():
outputs = self.model.generate(
**encoded,
max_new_tokens=parameters.get("max_new_tokens", 512),
temperature=parameters.get("temperature", 0.3),
do_sample=True,
repetition_penalty=1.2,
pad_token_id=self.tokenizer.eos_token_id,
)
generated = outputs[0][encoded["input_ids"].shape[1]:]
return [{"generated_text": self.tokenizer.decode(generated, skip_special_tokens=True)}]
|