glance-text / modeling_glance.py

Upload GLANCE

1469fed verified 25 days ago

16.3 kB

"""GLANCE: statement encoder. Dato un testo e uno o piu' statement, restituisce la probabilita' che ogni statement sia vero.
Questo file viaggia dentro il repository Hugging Face del modello, quindi basta:
    from transformers import AutoModel
    model = AutoModel.from_pretrained("utente/repo", trust_remote_code=True)
    model.predict("The vase is broken.", ["The vase is intact.", "It is sad."])  # -> array([0.02, 0.81])
    model.predict([("The vase is broken.", ["It is sad."]),                       # batch: un array per stato
                   ("I love it.", ["It is positive.", "It is sad."])])            # -> [array([0.81]), array([0.97, 0.04])]
Non e' un modello generativo: al posto di `generate` c'e' `predict`, che restituisce probabilita'.
Come funziona. Una sequenza contiene il testo (lo "stato") e alcuni statement:
    [CLS] testo ... [SEP] | [CLS] statement 1 [SEP] | [CLS] statement 2 [SEP] | ...
Una maschera di attenzione a blocchi fa si' che il testo veda solo se stesso e ogni statement veda solo il testo e se
stesso; le posizioni di ogni statement ripartono dalla fine del testo. Cosi' valutare molti statement insieme da' lo
stesso risultato che valutarli uno per uno, ma il testo viene letto una volta sola.
"""
import os  # percorsi

import numpy as np  # array
import torch  # tensori
import torch.nn as nn  # layer
import transformers  # per leggere la versione
from transformers import AutoModel, AutoTokenizer, PretrainedConfig, PreTrainedModel  # classi base di HF

_TF_MAJOR = int(transformers.__version__.split(".")[0])  # versione principale di transformers


class GlanceConfig(PretrainedConfig):  # configurazione del modello
    model_type = "glance"  # nome del tipo di modello

    # parametri con i default di training
    def __init__(self, pooling="cls", dropout=0.1, max_state_tokens=384, max_question_tokens=96,
                 max_pack_tokens=1024, max_questions_per_pack=8, **kwargs):
        super().__init__(**kwargs)  # parametri standard di HF
        self.pooling = pooling                          # "cls" o "mean": come si riassume ogni statement
        self.dropout = dropout
        self.max_state_tokens = max_state_tokens        # limiti usati in training: il testo oltre viene troncato
        self.max_question_tokens = max_question_tokens
        self.max_pack_tokens = max_pack_tokens          # lunghezza massima di una sequenza
        self.max_questions_per_pack = max_questions_per_pack


def _to_dtype(dtype):  # nome -> dtype di torch
    if dtype is None or dtype == "auto":  # nessuna scelta
        return None  # decide from_pretrained
    if isinstance(dtype, str):  # dato come stringa?
        # nome -> dtype
        return {"float16": torch.float16, "fp16": torch.float16, "bfloat16": torch.bfloat16, "bf16": torch.bfloat16,
                "float32": torch.float32, "fp32": torch.float32}[dtype]
    return dtype  # gia' un dtype


class Glance(PreTrainedModel):  # modello per AutoModel
    config_class = GlanceConfig  # classe di configurazione
    base_model_prefix = "backbone"  # nome del backbone

    def __init__(self, config, backbone=None, tokenizer=None):  # costruisce il modello
        super().__init__(config)  # init di HF
        self.backbone = backbone  # encoder
        # dimensione nascosta
        H = backbone.config.hidden_size if backbone is not None else getattr(config, "hidden_size", 768)
        # stessa struttura della testa di training (i Dropout non hanno pesi e in inferenza non fanno nulla)
        self.head = nn.Sequential(nn.Dropout(config.dropout), nn.Linear(H, H), nn.Tanh(),  # testa: H -> H -> 1 logit
                                  nn.Dropout(config.dropout), nn.Linear(H, 1))
        self.tokenizer = tokenizer  # tokenizer incluso

    # ------------------------------------------------------------------------------------------------------------
    # Caricamento
    # ------------------------------------------------------------------------------------------------------------
    @classmethod  # costruttore alternativo
    # carica da cartella o dal Hub
    def from_pretrained(cls, pretrained_model_name_or_path, *model_args, config=None, device=None, dtype=None,
                        **kwargs):
        """Carica encoder, testa e tokenizer da una cartella locale o da un repo del Hub.
        device: default "cuda" se disponibile, altrimenti "cpu". dtype: default float16 su GPU, float32 su CPU."""
        path = str(pretrained_model_name_or_path)  # percorso come stringa
        if not os.path.isdir(path):  # non e' una cartella locale?
            from huggingface_hub import snapshot_download  # import solo se serve
            # scarica il repo
            path = snapshot_download(path, revision=kwargs.get("_commit_hash") or kwargs.get("revision"),
                                     token=kwargs.get("token"), cache_dir=kwargs.get("cache_dir"),
                                     local_files_only=kwargs.get("local_files_only", False))
        if config is None:  # config non passata?
            config = GlanceConfig.from_pretrained(path)  # la legge dal repo

        backbone = AutoModel.from_pretrained(os.path.join(path, "backbone"),  # carica il backbone
                                             attn_implementation=kwargs.get("attn_implementation", "sdpa"))
        backbone.pooler = None                          # non serve: si usano i vettori dei singoli token
        # config=backbone.config: senza, AutoTokenizer leggerebbe il config.json della cartella (che punta a questo
        # codice) e chiederebbe di nuovo il permesso di eseguirlo
        tokenizer = AutoTokenizer.from_pretrained(path, config=backbone.config)  # carica il tokenizer
        model = cls(config, backbone=backbone, tokenizer=tokenizer)  # assembla il modello

        from safetensors.torch import load_file  # lettura dei pesi
        model.head.load_state_dict(load_file(os.path.join(path, "head.safetensors")))  # pesi della testa

        device = device or ("cuda" if torch.cuda.is_available() else "cpu")  # GPU se disponibile
        dtype = _to_dtype(dtype if dtype is not None else kwargs.get("torch_dtype"))  # precisione richiesta
        if dtype is None:  # nessuna scelta?
            dtype = torch.float16 if str(device).startswith("cuda") else torch.float32  # fp16 su GPU, fp32 su CPU
        return model.to(device=device, dtype=dtype).eval()  # sposta e mette in valutazione

    # ------------------------------------------------------------------------------------------------------------
    # Costruzione delle sequenze
    # ------------------------------------------------------------------------------------------------------------
    def _special_ids(self):  # id dei token speciali
        tok = self.tokenizer  # tokenizer
        cls_id = tok.cls_token_id if tok.cls_token_id is not None else tok.bos_token_id  # CLS o BOS
        sep_id = tok.sep_token_id if tok.sep_token_id is not None else tok.eos_token_id  # SEP o EOS
        pad_id = tok.pad_token_id if tok.pad_token_id is not None else 0  # padding, 0 se manca
        return cls_id, sep_id, pad_id  # i tre id

    def _tokenize(self, texts, max_tokens):  # testi -> id troncati
        # senza token speciali (li aggiungiamo noi), lasciando posto per [CLS] e [SEP]
        return self.tokenizer(list(texts), add_special_tokens=False, truncation=True,  # senza token speciali
                              max_length=max_tokens - 2)["input_ids"]

    def _build_packs(self, text_ids, statement_ids):  # testo + statement -> pack
        """Divide gli statement di un testo in una o piu' sequenze che rispettano i limiti di lunghezza.
        Restituisce (ids, segmenti, posizioni, dove sta il [CLS] di ogni statement, indici degli statement)."""
        cls_id, sep_id, _ = self._special_ids()  # token speciali
        c = self.config  # limiti
        S = [cls_id] + text_ids + [sep_id]  # testo: [CLS] ... [SEP]
        packs, current, used = [], [], 0  # pack finiti, in corso, token usati
        for k, q in enumerate(statement_ids):  # ogni statement
            Q = [cls_id] + q + [sep_id]  # [CLS] statement [SEP]
            # pack pieno?
            if current and (len(current) >= c.max_questions_per_pack or used + len(Q) > c.max_pack_tokens - len(S)):
                packs.append(current)  # lo chiude
                current, used = [], 0  # ne apre uno nuovo
            current.append((k, Q))  # aggiunge lo statement
            used += len(Q)  # aggiorna i token usati
        if current:  # resta un pack aperto?
            packs.append(current)  # lo chiude

        out = []  # pack assemblati
        for pack in packs:  # ogni pack
            ids, seg, pos, read, idx = list(S), [0] * len(S), list(range(len(S))), [], []  # parte dal testo
            for j, (k, Q) in enumerate(pack, start=1):  # ogni statement del pack
                read.append(len(ids))  # posizione del suo CLS
                idx.append(k)  # indice originale
                ids += Q  # accoda i token
                seg += [j] * len(Q)                     # segmento 0 = testo, j = statement j
                pos += range(len(S), len(S) + len(Q))   # ogni statement riparte dalla fine del testo
            out.append((ids, seg, pos, read, idx))  # pack finito
        return out  # lista di pack

    # ------------------------------------------------------------------------------------------------------------
    # Passaggio nel modello
    # ------------------------------------------------------------------------------------------------------------
    def forward(self, input_ids, segment_ids, position_ids, readout_index):  # tensori -> logit
        """Tensori gia' preparati -> logit [batch, statement]. Per l'uso normale c'e' `predict`."""
        L = input_ids.shape[1]  # lunghezza della sequenza
        # il token i puo' guardare il token j se stanno nello stesso segmento, oppure se j e' nel testo e i non e'
        # padding (segmento -1). Il testo non vede gli statement, gli statement non si vedono tra loro.
        s_i, s_j = segment_ids[:, :, None], segment_ids[:, None, :]  # segmenti di query e chiavi
        allow = ((s_j == s_i) | ((s_j == 0) & (s_i >= 0))) & (s_j >= 0)  # maschera a blocchi
        allow |= torch.eye(L, dtype=torch.bool, device=input_ids.device)  # ogni token vede se stesso

        emb = self.backbone.embeddings.position_embeddings  # embedding di posizione
        offset = emb.padding_idx + 1 if emb.padding_idx is not None else 0   # RoBERTa conta le posizioni da 2
        if getattr(self.backbone.config, "type_vocab_size", 1) > 1:  # BERT con due tipi di token?
            token_type = (segment_ids > 0).long()        # BERT: 0 per il testo, 1 per gli statement
        else:  # RoBERTa
            token_type = torch.zeros_like(input_ids)  # tutti zero

        if _TF_MAJOR < 5:  # transformers 4
            mask = allow.long()                          # transformers 4: maschera [B, L, L] di 0/1
        elif getattr(self.backbone.config, "_attn_implementation", "sdpa") == "eager":  # transformers 5 con eager
            neg = torch.finfo(self.backbone.dtype).min   # transformers 5 + eager: maschera additiva
            mask = torch.zeros(allow.shape, dtype=self.backbone.dtype, device=allow.device)  # 0 = visibile
            mask = mask.masked_fill(~allow, neg)[:, None]  # minimo = nascosto
        else:  # transformers 5 con sdpa
            mask = allow[:, None]                        # transformers 5 + sdpa: maschera booleana [B, 1, L, L]
        # hidden state di ogni token
        h = self.backbone(input_ids=input_ids, attention_mask=mask, token_type_ids=token_type,
                          position_ids=position_ids + offset).last_hidden_state

        if self.config.pooling == "mean":  # pooling medio
            Q = readout_index.shape[1]  # numero di statement
            k = torch.arange(1, Q + 1, device=h.device)[None, :, None]  # segmenti 1..Q
            w = (segment_ids[:, None, :] == k).to(h.dtype)  # token di ogni statement
            pooled = (w @ h) / w.sum(-1, keepdim=True).clamp(min=1)  # media dei token
        else:  # pooling sul CLS
            pooled = h.gather(1, readout_index[..., None].expand(-1, -1, h.shape[-1]))  # CLS di ogni statement
        return self.head(pooled).squeeze(-1)  # logit

    @torch.no_grad()  # niente gradienti
    def _run(self, packs):  # pack -> probabilita'
        _, _, pad_id = self._special_ids()  # id di padding
        B = len(packs)  # numero di pack
        L = max(len(p[0]) for p in packs)  # lunghezza massima
        K = max(len(p[3]) for p in packs)  # statement massimi
        ids = torch.full((B, L), pad_id)  # token, con padding
        seg = torch.full((B, L), -1)  # -1 = padding
        pos = torch.zeros((B, L), dtype=torch.long)  # posizioni
        read = torch.zeros((B, K), dtype=torch.long)  # indici dei CLS
        for b, (i, s, p, r, _) in enumerate(packs):  # copia ogni pack
            ids[b, :len(i)] = torch.tensor(i)  # token
            seg[b, :len(s)] = torch.tensor(s)  # segmenti
            pos[b, :len(p)] = torch.tensor(p)  # posizioni
            read[b, :len(r)] = torch.tensor(r)  # indici dei CLS
        dev = self.head[1].weight.device  # device del modello
        logits = self(ids.to(dev), seg.to(dev), pos.to(dev), read.to(dev))  # forward
        return torch.sigmoid(logits.float()).cpu().numpy()  # probabilita' su CPU

    # ------------------------------------------------------------------------------------------------------------
    # API semplice
    # ------------------------------------------------------------------------------------------------------------
    def predict(self, state, questions=None, batch_size=32):  # probabilita' che le domande siano vere
        """Due forme:
            model.predict("The vase is broken.", ["The sentence is negative.", "Something got damaged."])
            # -> array([0.93, 0.88])                     uno stato e le sue domande
            model.predict([("The vase is broken.", ["The sentence is negative."]),
                           {"state": "I love it.", "questions": ["It is positive.", "It is sad."]}])
            # -> [array([0.93]), array([0.97, 0.04])]    un batch: un array per stato
        Nel batch ogni elemento e' una coppia (stato, domande) oppure un record {"state": ..., "questions": ...}."""
        if questions is not None:  # uno stato e le sue domande
            return self._predict([(state, questions)], batch_size)[0]  # array per domanda
        if isinstance(state, str):  # uno stato senza domande
            raise ValueError("passa le domande: predict(stato, domande), oppure un batch di (stato, domande)")
        return self._predict(state, batch_size)  # batch

    def _predict(self, batch, batch_size):  # [(stato, domande), ...] -> [array, ...]
        pairs = [(x["state"], x["questions"]) if isinstance(x, dict) else x for x in batch]  # record -> coppia
        questions = [[qs] if isinstance(qs, str) else list(qs) for _, qs in pairs]  # domande, sempre liste
        c = self.config  # limiti
        state_ids = self._tokenize([str(s) for s, _ in pairs], c.max_state_tokens)  # tokenizza gli stati
        uniq = list(dict.fromkeys(q for qs in questions for q in qs))  # domande uniche
        q_ids = dict(zip(uniq, self._tokenize(uniq, c.max_question_tokens)))  # tokenizzate una volta sola

        # (stato, pack) per ogni pack
        jobs = [(r, pack) for r, (ids, qs) in enumerate(zip(state_ids, questions))
                for pack in self._build_packs(ids, [q_ids[q] for q in qs])]
        order = sorted(range(len(jobs)), key=lambda j: len(jobs[j][1][0]))   # lunghezze simili insieme
        out = [np.full(len(qs), np.nan) for qs in questions]  # risultati, NaN finche' non calcolati
        for start in range(0, len(order), batch_size):  # a blocchi di batch_size
            chunk = [jobs[j] for j in order[start: start + batch_size]]  # pack del blocco
            probs = self._run([pack for _, pack in chunk])  # probabilita' del blocco
            for b, (r, pack) in enumerate(chunk):  # ogni pack
                out[r][pack[4]] = probs[b, :len(pack[4])]  # rimette le probabilita' al loro posto
        return out  # una lista di array per stato

	"""GLANCE: statement encoder. Dato un testo e uno o piu' statement, restituisce la probabilita' che ogni statement sia vero.

	Questo file viaggia dentro il repository Hugging Face del modello, quindi basta:

	from transformers import AutoModel

	model = AutoModel.from_pretrained("utente/repo", trust_remote_code=True)
	model.predict("The vase is broken.", ["The vase is intact.", "It is sad."]) # -> array([0.02, 0.81])
	model.predict([("The vase is broken.", ["It is sad."]), # batch: un array per stato
	("I love it.", ["It is positive.", "It is sad."])]) # -> [array([0.81]), array([0.97, 0.04])]

	Non e' un modello generativo: al posto di `generate` c'e' `predict`, che restituisce probabilita'.

	Come funziona. Una sequenza contiene il testo (lo "stato") e alcuni statement:

	[CLS] testo ... [SEP] \| [CLS] statement 1 [SEP] \| [CLS] statement 2 [SEP] \| ...

	Una maschera di attenzione a blocchi fa si' che il testo veda solo se stesso e ogni statement veda solo il testo e se
	stesso; le posizioni di ogni statement ripartono dalla fine del testo. Cosi' valutare molti statement insieme da' lo
	stesso risultato che valutarli uno per uno, ma il testo viene letto una volta sola.
	"""
	import os # percorsi

	import numpy as np # array
	import torch # tensori
	import torch.nn as nn # layer
	import transformers # per leggere la versione
	from transformers import AutoModel, AutoTokenizer, PretrainedConfig, PreTrainedModel # classi base di HF

	_TF_MAJOR = int(transformers.__version__.split(".")[0]) # versione principale di transformers


	class GlanceConfig(PretrainedConfig): # configurazione del modello
	model_type = "glance" # nome del tipo di modello

	# parametri con i default di training
	def __init__(self, pooling="cls", dropout=0.1, max_state_tokens=384, max_question_tokens=96,
	max_pack_tokens=1024, max_questions_per_pack=8, **kwargs):
	super().__init__(**kwargs) # parametri standard di HF
	self.pooling = pooling # "cls" o "mean": come si riassume ogni statement
	self.dropout = dropout
	self.max_state_tokens = max_state_tokens # limiti usati in training: il testo oltre viene troncato
	self.max_question_tokens = max_question_tokens
	self.max_pack_tokens = max_pack_tokens # lunghezza massima di una sequenza
	self.max_questions_per_pack = max_questions_per_pack


	def _to_dtype(dtype): # nome -> dtype di torch
	if dtype is None or dtype == "auto": # nessuna scelta
	return None # decide from_pretrained
	if isinstance(dtype, str): # dato come stringa?
	# nome -> dtype
	return {"float16": torch.float16, "fp16": torch.float16, "bfloat16": torch.bfloat16, "bf16": torch.bfloat16,
	"float32": torch.float32, "fp32": torch.float32}[dtype]
	return dtype # gia' un dtype


	class Glance(PreTrainedModel): # modello per AutoModel
	config_class = GlanceConfig # classe di configurazione
	base_model_prefix = "backbone" # nome del backbone

	def __init__(self, config, backbone=None, tokenizer=None): # costruisce il modello
	super().__init__(config) # init di HF
	self.backbone = backbone # encoder
	# dimensione nascosta
	H = backbone.config.hidden_size if backbone is not None else getattr(config, "hidden_size", 768)
	# stessa struttura della testa di training (i Dropout non hanno pesi e in inferenza non fanno nulla)
	self.head = nn.Sequential(nn.Dropout(config.dropout), nn.Linear(H, H), nn.Tanh(), # testa: H -> H -> 1 logit
	nn.Dropout(config.dropout), nn.Linear(H, 1))
	self.tokenizer = tokenizer # tokenizer incluso

	# ------------------------------------------------------------------------------------------------------------
	# Caricamento
	# ------------------------------------------------------------------------------------------------------------
	@classmethod # costruttore alternativo
	# carica da cartella o dal Hub
	def from_pretrained(cls, pretrained_model_name_or_path, *model_args, config=None, device=None, dtype=None,
	**kwargs):
	"""Carica encoder, testa e tokenizer da una cartella locale o da un repo del Hub.
	device: default "cuda" se disponibile, altrimenti "cpu". dtype: default float16 su GPU, float32 su CPU."""
	path = str(pretrained_model_name_or_path) # percorso come stringa
	if not os.path.isdir(path): # non e' una cartella locale?
	from huggingface_hub import snapshot_download # import solo se serve
	# scarica il repo
	path = snapshot_download(path, revision=kwargs.get("_commit_hash") or kwargs.get("revision"),
	token=kwargs.get("token"), cache_dir=kwargs.get("cache_dir"),
	local_files_only=kwargs.get("local_files_only", False))
	if config is None: # config non passata?
	config = GlanceConfig.from_pretrained(path) # la legge dal repo

	backbone = AutoModel.from_pretrained(os.path.join(path, "backbone"), # carica il backbone
	attn_implementation=kwargs.get("attn_implementation", "sdpa"))
	backbone.pooler = None # non serve: si usano i vettori dei singoli token
	# config=backbone.config: senza, AutoTokenizer leggerebbe il config.json della cartella (che punta a questo
	# codice) e chiederebbe di nuovo il permesso di eseguirlo
	tokenizer = AutoTokenizer.from_pretrained(path, config=backbone.config) # carica il tokenizer
	model = cls(config, backbone=backbone, tokenizer=tokenizer) # assembla il modello

	from safetensors.torch import load_file # lettura dei pesi
	model.head.load_state_dict(load_file(os.path.join(path, "head.safetensors"))) # pesi della testa

	device = device or ("cuda" if torch.cuda.is_available() else "cpu") # GPU se disponibile
	dtype = _to_dtype(dtype if dtype is not None else kwargs.get("torch_dtype")) # precisione richiesta
	if dtype is None: # nessuna scelta?
	dtype = torch.float16 if str(device).startswith("cuda") else torch.float32 # fp16 su GPU, fp32 su CPU
	return model.to(device=device, dtype=dtype).eval() # sposta e mette in valutazione

	# ------------------------------------------------------------------------------------------------------------
	# Costruzione delle sequenze
	# ------------------------------------------------------------------------------------------------------------
	def _special_ids(self): # id dei token speciali
	tok = self.tokenizer # tokenizer
	cls_id = tok.cls_token_id if tok.cls_token_id is not None else tok.bos_token_id # CLS o BOS
	sep_id = tok.sep_token_id if tok.sep_token_id is not None else tok.eos_token_id # SEP o EOS
	pad_id = tok.pad_token_id if tok.pad_token_id is not None else 0 # padding, 0 se manca
	return cls_id, sep_id, pad_id # i tre id

	def _tokenize(self, texts, max_tokens): # testi -> id troncati
	# senza token speciali (li aggiungiamo noi), lasciando posto per [CLS] e [SEP]
	return self.tokenizer(list(texts), add_special_tokens=False, truncation=True, # senza token speciali
	max_length=max_tokens - 2)["input_ids"]

	def _build_packs(self, text_ids, statement_ids): # testo + statement -> pack
	"""Divide gli statement di un testo in una o piu' sequenze che rispettano i limiti di lunghezza.
	Restituisce (ids, segmenti, posizioni, dove sta il [CLS] di ogni statement, indici degli statement)."""
	cls_id, sep_id, _ = self._special_ids() # token speciali
	c = self.config # limiti
	S = [cls_id] + text_ids + [sep_id] # testo: [CLS] ... [SEP]
	packs, current, used = [], [], 0 # pack finiti, in corso, token usati
	for k, q in enumerate(statement_ids): # ogni statement
	Q = [cls_id] + q + [sep_id] # [CLS] statement [SEP]
	# pack pieno?
	if current and (len(current) >= c.max_questions_per_pack or used + len(Q) > c.max_pack_tokens - len(S)):
	packs.append(current) # lo chiude
	current, used = [], 0 # ne apre uno nuovo
	current.append((k, Q)) # aggiunge lo statement
	used += len(Q) # aggiorna i token usati
	if current: # resta un pack aperto?
	packs.append(current) # lo chiude

	out = [] # pack assemblati
	for pack in packs: # ogni pack
	ids, seg, pos, read, idx = list(S), [0] * len(S), list(range(len(S))), [], [] # parte dal testo
	for j, (k, Q) in enumerate(pack, start=1): # ogni statement del pack
	read.append(len(ids)) # posizione del suo CLS
	idx.append(k) # indice originale
	ids += Q # accoda i token
	seg += [j] * len(Q) # segmento 0 = testo, j = statement j
	pos += range(len(S), len(S) + len(Q)) # ogni statement riparte dalla fine del testo
	out.append((ids, seg, pos, read, idx)) # pack finito
	return out # lista di pack

	# ------------------------------------------------------------------------------------------------------------
	# Passaggio nel modello
	# ------------------------------------------------------------------------------------------------------------
	def forward(self, input_ids, segment_ids, position_ids, readout_index): # tensori -> logit
	"""Tensori gia' preparati -> logit [batch, statement]. Per l'uso normale c'e' `predict`."""
	L = input_ids.shape[1] # lunghezza della sequenza
	# il token i puo' guardare il token j se stanno nello stesso segmento, oppure se j e' nel testo e i non e'
	# padding (segmento -1). Il testo non vede gli statement, gli statement non si vedono tra loro.
	s_i, s_j = segment_ids[:, :, None], segment_ids[:, None, :] # segmenti di query e chiavi
	allow = ((s_j == s_i) \| ((s_j == 0) & (s_i >= 0))) & (s_j >= 0) # maschera a blocchi
	allow \|= torch.eye(L, dtype=torch.bool, device=input_ids.device) # ogni token vede se stesso

	emb = self.backbone.embeddings.position_embeddings # embedding di posizione
	offset = emb.padding_idx + 1 if emb.padding_idx is not None else 0 # RoBERTa conta le posizioni da 2
	if getattr(self.backbone.config, "type_vocab_size", 1) > 1: # BERT con due tipi di token?
	token_type = (segment_ids > 0).long() # BERT: 0 per il testo, 1 per gli statement
	else: # RoBERTa
	token_type = torch.zeros_like(input_ids) # tutti zero

	if _TF_MAJOR < 5: # transformers 4
	mask = allow.long() # transformers 4: maschera [B, L, L] di 0/1
	elif getattr(self.backbone.config, "_attn_implementation", "sdpa") == "eager": # transformers 5 con eager
	neg = torch.finfo(self.backbone.dtype).min # transformers 5 + eager: maschera additiva
	mask = torch.zeros(allow.shape, dtype=self.backbone.dtype, device=allow.device) # 0 = visibile
	mask = mask.masked_fill(~allow, neg)[:, None] # minimo = nascosto
	else: # transformers 5 con sdpa
	mask = allow[:, None] # transformers 5 + sdpa: maschera booleana [B, 1, L, L]
	# hidden state di ogni token
	h = self.backbone(input_ids=input_ids, attention_mask=mask, token_type_ids=token_type,
	position_ids=position_ids + offset).last_hidden_state

	if self.config.pooling == "mean": # pooling medio
	Q = readout_index.shape[1] # numero di statement
	k = torch.arange(1, Q + 1, device=h.device)[None, :, None] # segmenti 1..Q
	w = (segment_ids[:, None, :] == k).to(h.dtype) # token di ogni statement
	pooled = (w @ h) / w.sum(-1, keepdim=True).clamp(min=1) # media dei token
	else: # pooling sul CLS
	pooled = h.gather(1, readout_index[..., None].expand(-1, -1, h.shape[-1])) # CLS di ogni statement
	return self.head(pooled).squeeze(-1) # logit

	@torch.no_grad() # niente gradienti
	def _run(self, packs): # pack -> probabilita'
	_, _, pad_id = self._special_ids() # id di padding
	B = len(packs) # numero di pack
	L = max(len(p[0]) for p in packs) # lunghezza massima
	K = max(len(p[3]) for p in packs) # statement massimi
	ids = torch.full((B, L), pad_id) # token, con padding
	seg = torch.full((B, L), -1) # -1 = padding
	pos = torch.zeros((B, L), dtype=torch.long) # posizioni
	read = torch.zeros((B, K), dtype=torch.long) # indici dei CLS
	for b, (i, s, p, r, _) in enumerate(packs): # copia ogni pack
	ids[b, :len(i)] = torch.tensor(i) # token
	seg[b, :len(s)] = torch.tensor(s) # segmenti
	pos[b, :len(p)] = torch.tensor(p) # posizioni
	read[b, :len(r)] = torch.tensor(r) # indici dei CLS
	dev = self.head[1].weight.device # device del modello
	logits = self(ids.to(dev), seg.to(dev), pos.to(dev), read.to(dev)) # forward
	return torch.sigmoid(logits.float()).cpu().numpy() # probabilita' su CPU

	# ------------------------------------------------------------------------------------------------------------
	# API semplice
	# ------------------------------------------------------------------------------------------------------------
	def predict(self, state, questions=None, batch_size=32): # probabilita' che le domande siano vere
	"""Due forme:

	model.predict("The vase is broken.", ["The sentence is negative.", "Something got damaged."])
	# -> array([0.93, 0.88]) uno stato e le sue domande

	model.predict([("The vase is broken.", ["The sentence is negative."]),
	{"state": "I love it.", "questions": ["It is positive.", "It is sad."]}])
	# -> [array([0.93]), array([0.97, 0.04])] un batch: un array per stato

	Nel batch ogni elemento e' una coppia (stato, domande) oppure un record {"state": ..., "questions": ...}."""
	if questions is not None: # uno stato e le sue domande
	return self._predict([(state, questions)], batch_size)[0] # array per domanda
	if isinstance(state, str): # uno stato senza domande
	raise ValueError("passa le domande: predict(stato, domande), oppure un batch di (stato, domande)")
	return self._predict(state, batch_size) # batch

	def _predict(self, batch, batch_size): # [(stato, domande), ...] -> [array, ...]
	pairs = [(x["state"], x["questions"]) if isinstance(x, dict) else x for x in batch] # record -> coppia
	questions = [[qs] if isinstance(qs, str) else list(qs) for _, qs in pairs] # domande, sempre liste
	c = self.config # limiti
	state_ids = self._tokenize([str(s) for s, _ in pairs], c.max_state_tokens) # tokenizza gli stati
	uniq = list(dict.fromkeys(q for qs in questions for q in qs)) # domande uniche
	q_ids = dict(zip(uniq, self._tokenize(uniq, c.max_question_tokens))) # tokenizzate una volta sola

	# (stato, pack) per ogni pack
	jobs = [(r, pack) for r, (ids, qs) in enumerate(zip(state_ids, questions))
	for pack in self._build_packs(ids, [q_ids[q] for q in qs])]
	order = sorted(range(len(jobs)), key=lambda j: len(jobs[j][1][0])) # lunghezze simili insieme
	out = [np.full(len(qs), np.nan) for qs in questions] # risultati, NaN finche' non calcolati
	for start in range(0, len(order), batch_size): # a blocchi di batch_size
	chunk = [jobs[j] for j in order[start: start + batch_size]] # pack del blocco
	probs = self._run([pack for _, pack in chunk]) # probabilita' del blocco
	for b, (r, pack) in enumerate(chunk): # ogni pack
	out[r][pack[4]] = probs[b, :len(pack[4])] # rimette le probabilita' al loro posto
	return out # una lista di array per stato