Release

659438c 4 days ago

7.07 kB

"""Calibrated typed decisions over audio, images or text by Abacus.AI: one forward pass, a probability per option, no generated text.
    from smaug_decide import SmaugDecide
    model = SmaugDecide.from_pretrained('abacusai/Smaug-Decisions')
    model('clip.wav', 'Has the speaker finished their turn?')                      # -> {'Yes': 0.93, 'No': 0.07}
    model('clip.wav', 'Which genre is playing?', ['rock', 'jazz', 'classical'])    # -> one-answer: probabilities sum to 1
    model('clip.wav', 'Which sounds are present?', ['speech', 'music', 'alarm'], multi=True)   # -> select-all: one probability each
    model('photo.jpg', 'Is there a person in the picture?')                        # images and plain text work the same way
    model('The order never arrived.', 'How upset is the customer?', ['calm', 'annoyed', 'angry'])
The readout is the model's own LM-head logits for the option letters A..T at the answer position (softcapped), softmax for one-answer
questions and sigmoid(z - b) for select-all questions; `b` is the learned select-all threshold stored in readout.json / head.pt.
Runs on NVIDIA GPUs (bfloat16), Apple Silicon (MPS) and CPU; the device is picked automatically.
"""
import json
import os

import numpy as np
import soundfile as sf
import torch
from huggingface_hub import snapshot_download
from PIL import Image
from transformers import AutoModelForMultimodalLM, AutoProcessor


LETTERS = 'ABCDEFGHIJKLMNOPQRST'
SAMPLE_RATE, MAX_AUDIO_SECONDS = 16000, 30.0            # the model hears 16 kHz mono; longer clips keep their last 30 s
MAX_TOKENS, AUDIO_TOKENS_PER_SECOND, IMAGE_TOKENS = 3072, 25, 560   # images at 560 soft tokens, as trained and evaluated
PREAMBLE = ('You make decisions for apps. Read the rules, the context and the input, then answer the question by choosing '
            'from the lettered options.')
IMAGE_SUFFIXES = ('.jpg', '.jpeg', '.png', '.webp', '.gif', '.bmp')
AUDIO_SUFFIXES = ('.wav', '.flac', '.ogg', '.mp3', '.m4a', '.aiff', '.aif')
WEIGHT_FILES = ['*.json', '*.jinja', 'model*.safetensors', 'head.pt']   # the checkpoint at the repo root


def trim(s, n):
    if len(s) <= n:
        return s
    return s[:int(n * 0.7)] + ' ... ' + s[-int(n * 0.3):]


def render(question, options, descriptions, multi, media_tag, rules, context, text_input, char_budget):
    """The training prompt, verbatim: preamble, optional rules / context, the input, the lettered options, the answer instruction."""
    parts = [PREAMBLE]
    if rules:
        parts.append(f'Rules: {trim(rules, char_budget // 2)}')
    if context:
        parts.append(f'Context: {trim(context, char_budget)}')
    if media_tag:
        parts.append(f'Input: {media_tag}')
    elif text_input:
        parts.append(f'Input: {trim(text_input, char_budget)}')
    lines = [f'Question: {question}' + (' Select all that apply.' if multi else '')]
    for k, (opt, desc) in enumerate(zip(options, descriptions)):
        lines.append(f'{LETTERS[k]}) {opt}' + (f' — {desc}' if desc else ''))
    parts.append('\n'.join(lines))
    parts.append('Answer with the letters of all options that apply.' if multi else 'Answer with the letter of the best option.')
    return '<bos><|turn>user\n' + '\n'.join(parts) + '<turn|>\n<|turn>model\n'


def load_audio(x):
    """A path, or a (samples, sample_rate) pair, or a 1-D float array at 16 kHz -> mono float32 at 16 kHz, last 30 s."""
    if isinstance(x, str):
        w, sr = sf.read(x, dtype='float32')
    elif isinstance(x, tuple):
        w, sr = x
    else:
        w, sr = x, SAMPLE_RATE
    w = np.asarray(w, dtype=np.float32)
    if w.ndim > 1:
        w = w.mean(1)
    if sr != SAMPLE_RATE:   # linear resampling is enough for a reference script; use a proper resampler for production audio
        n = int(round(len(w) * SAMPLE_RATE / sr))
        w = np.interp(np.linspace(0.0, len(w) - 1, n), np.arange(len(w)), w).astype(np.float32)
    return w[-int(MAX_AUDIO_SECONDS * SAMPLE_RATE):]


class SmaugDecide:
    def __init__(self, model, processor, readout, device):
        self.model, self.proc, self.device = model.eval(), processor, device
        tok = processor.tokenizer
        self.letter_ids = torch.tensor([tok(L, add_special_tokens=False)['input_ids'][0] for L in LETTERS], device=device)
        self.cap, self.b = readout.get('final_logit_softcapping'), float(readout['multi_bias_b'])

    @classmethod
    def from_pretrained(cls, name_or_path, device=None):
        path = name_or_path if os.path.isdir(name_or_path) else snapshot_download(name_or_path, allow_patterns=WEIGHT_FILES)
        device = device or ('cuda' if torch.cuda.is_available() else 'mps' if torch.backends.mps.is_available() else 'cpu')
        model = AutoModelForMultimodalLM.from_pretrained(path, dtype=torch.bfloat16).to(device)
        readout = json.load(open(os.path.join(path, 'readout.json')))
        return cls(model, AutoProcessor.from_pretrained(path), readout, device)

    @torch.no_grad()
    def __call__(self, inp, question, options=('Yes', 'No'), descriptions=None, multi=False, rules=None, context=None):
        """inp: an audio path / (samples, sr) / 16 kHz array, an image path / PIL image, or a plain string. Returns {option: probability}."""
        options = list(options)
        assert 2 <= len(options) <= len(LETTERS), f'2 to {len(LETTERS)} options'
        descriptions = list(descriptions) if descriptions else [None] * len(options)
        audio = image = text_input = None
        if isinstance(inp, Image.Image) or (isinstance(inp, str) and inp.lower().endswith(IMAGE_SUFFIXES)):
            image = (inp if isinstance(inp, Image.Image) else Image.open(inp)).convert('RGB')
        elif isinstance(inp, str) and not inp.lower().endswith(AUDIO_SUFFIXES) and not os.path.exists(inp):
            text_input = inp
        else:
            audio = load_audio(inp)
        media_tokens = len(audio) / SAMPLE_RATE * AUDIO_TOKENS_PER_SECOND + 2 if audio is not None else (IMAGE_TOKENS + 2 if image is not None else 0)
        tag = '<|audio|>' if audio is not None else '<|image|>' if image is not None else None
        budget = int(max(400, (MAX_TOKENS - media_tokens - 80 - 10 * len(options)) * 3.5))
        text = render(question, options, descriptions, multi, tag, rules, context, text_input, budget)
        kw = {}
        if audio is not None:
            kw['audio'] = [audio]
        if image is not None:
            kw['images'] = [image]; kw['images_kwargs'] = {'max_soft_tokens': IMAGE_TOKENS}
        enc = self.proc(text=[text], return_tensors='pt', add_special_tokens=False, **kw).to(self.device)
        h = self.model.model(**enc, use_cache=False).last_hidden_state[0, -1].float()          # hidden state at the answer position
        z = h @ self.model.lm_head.weight[self.letter_ids[:len(options)]].float().T          # LM-head logits of the letters only
        if self.cap:
            z = torch.tanh(z / self.cap) * self.cap
        p = torch.sigmoid(z - self.b) if multi else torch.softmax(z, -1)
        return dict(zip(options, p.tolist()))

	"""Calibrated typed decisions over audio, images or text by Abacus.AI: one forward pass, a probability per option, no generated text.

	from smaug_decide import SmaugDecide
	model = SmaugDecide.from_pretrained('abacusai/Smaug-Decisions')
	model('clip.wav', 'Has the speaker finished their turn?') # -> {'Yes': 0.93, 'No': 0.07}
	model('clip.wav', 'Which genre is playing?', ['rock', 'jazz', 'classical']) # -> one-answer: probabilities sum to 1
	model('clip.wav', 'Which sounds are present?', ['speech', 'music', 'alarm'], multi=True) # -> select-all: one probability each
	model('photo.jpg', 'Is there a person in the picture?') # images and plain text work the same way
	model('The order never arrived.', 'How upset is the customer?', ['calm', 'annoyed', 'angry'])

	The readout is the model's own LM-head logits for the option letters A..T at the answer position (softcapped), softmax for one-answer
	questions and sigmoid(z - b) for select-all questions; `b` is the learned select-all threshold stored in readout.json / head.pt.
	Runs on NVIDIA GPUs (bfloat16), Apple Silicon (MPS) and CPU; the device is picked automatically.
	"""
	import json
	import os

	import numpy as np
	import soundfile as sf
	import torch
	from huggingface_hub import snapshot_download
	from PIL import Image
	from transformers import AutoModelForMultimodalLM, AutoProcessor


	LETTERS = 'ABCDEFGHIJKLMNOPQRST'
	SAMPLE_RATE, MAX_AUDIO_SECONDS = 16000, 30.0 # the model hears 16 kHz mono; longer clips keep their last 30 s
	MAX_TOKENS, AUDIO_TOKENS_PER_SECOND, IMAGE_TOKENS = 3072, 25, 560 # images at 560 soft tokens, as trained and evaluated
	PREAMBLE = ('You make decisions for apps. Read the rules, the context and the input, then answer the question by choosing '
	'from the lettered options.')
	IMAGE_SUFFIXES = ('.jpg', '.jpeg', '.png', '.webp', '.gif', '.bmp')
	AUDIO_SUFFIXES = ('.wav', '.flac', '.ogg', '.mp3', '.m4a', '.aiff', '.aif')
	WEIGHT_FILES = ['.json', '.jinja', 'model*.safetensors', 'head.pt'] # the checkpoint at the repo root


	def trim(s, n):
	if len(s) <= n:
	return s
	return s[:int(n * 0.7)] + ' ... ' + s[-int(n * 0.3):]


	def render(question, options, descriptions, multi, media_tag, rules, context, text_input, char_budget):
	"""The training prompt, verbatim: preamble, optional rules / context, the input, the lettered options, the answer instruction."""
	parts = [PREAMBLE]
	if rules:
	parts.append(f'Rules: {trim(rules, char_budget // 2)}')
	if context:
	parts.append(f'Context: {trim(context, char_budget)}')
	if media_tag:
	parts.append(f'Input: {media_tag}')
	elif text_input:
	parts.append(f'Input: {trim(text_input, char_budget)}')
	lines = [f'Question: {question}' + (' Select all that apply.' if multi else '')]
	for k, (opt, desc) in enumerate(zip(options, descriptions)):
	lines.append(f'{LETTERS[k]}) {opt}' + (f' — {desc}' if desc else ''))
	parts.append('\n'.join(lines))
	parts.append('Answer with the letters of all options that apply.' if multi else 'Answer with the letter of the best option.')
	return '<bos><\|turn>user\n' + '\n'.join(parts) + '<turn\|>\n<\|turn>model\n'


	def load_audio(x):
	"""A path, or a (samples, sample_rate) pair, or a 1-D float array at 16 kHz -> mono float32 at 16 kHz, last 30 s."""
	if isinstance(x, str):
	w, sr = sf.read(x, dtype='float32')
	elif isinstance(x, tuple):
	w, sr = x
	else:
	w, sr = x, SAMPLE_RATE
	w = np.asarray(w, dtype=np.float32)
	if w.ndim > 1:
	w = w.mean(1)
	if sr != SAMPLE_RATE: # linear resampling is enough for a reference script; use a proper resampler for production audio
	n = int(round(len(w) * SAMPLE_RATE / sr))
	w = np.interp(np.linspace(0.0, len(w) - 1, n), np.arange(len(w)), w).astype(np.float32)
	return w[-int(MAX_AUDIO_SECONDS * SAMPLE_RATE):]


	class SmaugDecide:
	def __init__(self, model, processor, readout, device):
	self.model, self.proc, self.device = model.eval(), processor, device
	tok = processor.tokenizer
	self.letter_ids = torch.tensor([tok(L, add_special_tokens=False)['input_ids'][0] for L in LETTERS], device=device)
	self.cap, self.b = readout.get('final_logit_softcapping'), float(readout['multi_bias_b'])

	@classmethod
	def from_pretrained(cls, name_or_path, device=None):
	path = name_or_path if os.path.isdir(name_or_path) else snapshot_download(name_or_path, allow_patterns=WEIGHT_FILES)
	device = device or ('cuda' if torch.cuda.is_available() else 'mps' if torch.backends.mps.is_available() else 'cpu')
	model = AutoModelForMultimodalLM.from_pretrained(path, dtype=torch.bfloat16).to(device)
	readout = json.load(open(os.path.join(path, 'readout.json')))
	return cls(model, AutoProcessor.from_pretrained(path), readout, device)

	@torch.no_grad()
	def __call__(self, inp, question, options=('Yes', 'No'), descriptions=None, multi=False, rules=None, context=None):
	"""inp: an audio path / (samples, sr) / 16 kHz array, an image path / PIL image, or a plain string. Returns {option: probability}."""
	options = list(options)
	assert 2 <= len(options) <= len(LETTERS), f'2 to {len(LETTERS)} options'
	descriptions = list(descriptions) if descriptions else [None] * len(options)
	audio = image = text_input = None
	if isinstance(inp, Image.Image) or (isinstance(inp, str) and inp.lower().endswith(IMAGE_SUFFIXES)):
	image = (inp if isinstance(inp, Image.Image) else Image.open(inp)).convert('RGB')
	elif isinstance(inp, str) and not inp.lower().endswith(AUDIO_SUFFIXES) and not os.path.exists(inp):
	text_input = inp
	else:
	audio = load_audio(inp)
	media_tokens = len(audio) / SAMPLE_RATE * AUDIO_TOKENS_PER_SECOND + 2 if audio is not None else (IMAGE_TOKENS + 2 if image is not None else 0)
	tag = '<\|audio\|>' if audio is not None else '<\|image\|>' if image is not None else None
	budget = int(max(400, (MAX_TOKENS - media_tokens - 80 - 10 * len(options)) * 3.5))
	text = render(question, options, descriptions, multi, tag, rules, context, text_input, budget)
	kw = {}
	if audio is not None:
	kw['audio'] = [audio]
	if image is not None:
	kw['images'] = [image]; kw['images_kwargs'] = {'max_soft_tokens': IMAGE_TOKENS}
	enc = self.proc(text=[text], return_tensors='pt', add_special_tokens=False, **kw).to(self.device)
	h = self.model.model(**enc, use_cache=False).last_hidden_state[0, -1].float() # hidden state at the answer position
	z = h @ self.model.lm_head.weight[self.letter_ids[:len(options)]].float().T # LM-head logits of the letters only
	if self.cap:
	z = torch.tanh(z / self.cap) * self.cap
	p = torch.sigmoid(z - self.b) if multi else torch.softmax(z, -1)
	return dict(zip(options, p.tolist()))