Views
No views yet
jua-4B-mixed is a Brazilian Portuguese legal embedding model based on Qwen/Qwen3-Embedding-4B. It was adapted with a mixed supervision regime that combines legal-domain supervision with broader question-passage supervision, and is intended for heterogeneous retrieval settings where legal specialization and broader semantic robustness are both important.mixed condition discussed in the paper.Qwen/Qwen3-Embedding-4BJUÁ-Juris training pairsSQuAD-pt as a broader question-passage supervision sourceufca-llms/jua-4B-legal-only1# Requires transformers>=4.51.0
2# Requires sentence-transformers>=2.7.0
3
4from sentence_transformers import SentenceTransformer
5
6model = SentenceTransformer("ufca-llms/jua-4B-mixed")
7
8queries = [
9 "Instruct: Given a Brazilian legal search query, retrieve relevant legal passages or documents.\nQuery: aposentadoria por pensão estatutária",
10 "Instruct: Given a Brazilian legal search query, retrieve relevant legal passages or documents.\nQuery: por que dividir um país em estados?",
11]
12
13documents = [
14 "O art. 5º da Lei 9.717/1998 trata do regime previdenciário dos servidores públicos.",
15 "A divisão de um país em estados distribui competências administrativas e políticas em sistemas federativos.",
16]
17
18query_embeddings = model.encode(queries)
19document_embeddings = model.encode(documents)
20
21similarity = model.similarity(query_embeddings, document_embeddings)
22print(similarity)1# Requires transformers>=4.51.0
2
3import torch
4import torch.nn.functional as F
5
6from torch import Tensor
7from transformers import AutoModel, AutoTokenizer
8
9
10def last_token_pool(last_hidden_states: Tensor, attention_mask: Tensor) -> Tensor:
11 left_padding = (attention_mask[:, -1].sum() == attention_mask.shape[0])
12 if left_padding:
13 return last_hidden_states[:, -1]
14 sequence_lengths = attention_mask.sum(dim=1) - 1
15 batch_size = last_hidden_states.shape[0]
16 return last_hidden_states[
17 torch.arange(batch_size, device=last_hidden_states.device),
18 sequence_lengths,
19 ]
20
21
22def get_detailed_instruct(task_description: str, query: str) -> str:
23 return f"Instruct: {task_description}\nQuery: {query}"
24
25
26task = "Given a Brazilian legal search query, retrieve relevant legal passages or documents."
27queries = [
28 get_detailed_instruct(task, "aposentadoria por pensão estatutária"),
29 get_detailed_instruct(task, "por que dividir um país em estados?"),
30]
31
32documents = [
33 "O art. 5º da Lei 9.717/1998 trata do regime previdenciário dos servidores públicos.",
34 "A divisão de um país em estados distribui competências administrativas e políticas em sistemas federativos.",
35]
36
37input_texts = queries + documents
38
39tokenizer = AutoTokenizer.from_pretrained(
40 "ufca-llms/jua-4B-mixed",
41 padding_side="left",
42)
43model = AutoModel.from_pretrained("ufca-llms/jua-4B-mixed")
44
45batch_dict = tokenizer(
46 input_texts,
47 padding=True,
48 truncation=True,
49 max_length=8192,
50 return_tensors="pt",
51)
52batch_dict.to(model.device)
53
54outputs = model(**batch_dict)
55embeddings = last_token_pool(outputs.last_hidden_state, batch_dict["attention_mask"])
56embeddings = F.normalize(embeddings, p=2, dim=1)
57
58scores = embeddings[: len(queries)] @ embeddings[len(queries) :].T
59print(scores.tolist())mixed results reported in the paper over the five legal datasets in the JUÁ evaluation environment plus Quati.| Dataset | NDCG@10 | MRR@10 | MAP@10 |
|---|---|---|---|
| JUÁ-Juris | 0.290 | 0.230 | 0.231 |
| JurisTCU | 0.363 | 0.641 | 0.170 |
| NormasTCU | 0.305 | 0.474 | 0.184 |
| Ulysses-RFCorpus | 0.441 | 0.624 | 0.315 |
| BR-TaxQA-R | 0.777 | 0.800 | 0.701 |
| Quati | 0.503 | 0.799 | 0.247 |
| Average | 0.447 | 0.595 | 0.308 |
JUÁ-Juris, JurisTCU, NormasTCU, and BR-TaxQA-R), this model obtains:NDCG@10: 0.434MRR@10: 0.536MAP@10: 0.3211@misc{pereira2026domainadaptivedenseretrievalbrazilian,
2 title={Domain-Adaptive Dense Retrieval for Brazilian Legal Search},
3 author={Jayr Pereira and Roberto Lotufo and Luiz Bonifacio},
4 year={2026},
5 eprint={2605.04005},
6 archivePrefix={arXiv},
7 primaryClass={cs.IR},
8 url={https://arxiv.org/abs/2605.04005},
9}