from transformers import AutoModel, AutoTokenizer

model_name = "jinaai/jina-embeddings-v2-base-en"
tokenizer = AutoTokenizer.from_pretrained(model_name, trust_remote_code=True)
model = AutoModel.from_pretrained(model_name, trust_remote_code=True)

TypeChunks = list[str]
TypeSpans = list[tuple[int, int]]

def chunk_by_sentences(text: str) -> tuple[TypeChunks, TypeSpans]:
    """
    Input:
        text (str): "<cls> / 我輩 / は / 猫 / で / ある / 。/ 名前 / は / まだ / ない / 。/ <sep>"
    Returns:
        chunks (TypeChunks): ["我輩は猫である。", "名前はまだない。"]
        spans (TypeSpans): [(1, 7), (7, 12)]
    """

def late_chunking(
    model_output: 'BatchEncoding', span_annotation: list, max_length=None
):
    token_embeddings = model_output[0]
    outputs = []
    for embeddings, annotations in zip(token_embeddings, span_annotation):
        if (
            max_length is not None
        ):  # remove annotations which go bejond the max-length of the model
            annotations = [
                (start, min(end, max_length - 1))
                for (start, end) in annotations
                if start < (max_length - 1)
            ]
        pooled_embeddings = [
            embeddings[start:end].sum(dim=0) / (end - start)
            for start, end in annotations
            if (end - start) >= 1
        ]
        pooled_embeddings = [
            embedding.detach().cpu().numpy() for embedding in pooled_embeddings
        ]
        outputs.append(pooled_embeddings)
    return outputs

>>> inputs = tokenizer(input_text, return_tensors='pt')
>>> model_output = model(**inputs)
>>> embeddings_chunks = late_chunking(model_output, [span_annotations])[0]
>>> embeddings_chunks[0].shape
(768,)

# -*- coding: utf_8 -*-
from concurrent.futures import ThreadPoolExecutor

import numpy as np
import numpy.typing as npt
from openai import OpenAI


def generate_embeddings(
    batch_text: list[str],
    model_name="text-embedding-3-large",
    client: OpenAI = OpenAI(),
    dimension: int = 3072
) -> npt.NDArray[np.float64]:
    response = client.embeddings.create(
        input=batch_text,
        model=model_name,
        dimensions=dimension
    )
    return np.array([data.embedding for data in response.data])


def generate_matryoshka_embeddings(
    batch_text: list[str],
    model_name="text-embedding-3-large",
    client: OpenAI = OpenAI(),
    dimensions: list[int] = [64, 128, 256, 512, 1024]
) -> dict[int, npt.NDArray[np.float64]]:
    with ThreadPoolExecutor(max_workers=8) as executor:
        futures = []
        for dimension in dimensions:
            futures.append(
                executor.submit(
                    generate_embeddings,
                    batch_text,
                    model_name=model_name,
                    client=client,
                    dimension=dimension
                )
            )
    matryoshka_embeddings = dict()
    for future in futures:
        result = future.result()
        matryoshka_embeddings[result.shape[1]] = result
    return matryoshka_embeddings

>>> batch_text = ["hello world"]
>>> client = OpenAI()
>>> dimensions = [64, 128, 256, 512, 1024]
>>> embeddings = generate_matryoshka_embeddings(batch_text, client=client, dimensions=dimensions)
>>> embeddings[64].shape
(1, 64)

あなたは科学者のためのアシスタントです。ユーザーからの質問が提供されます。
あなたの目標は、科学者が文献に質問を投げかけ、**それに答える準備をするための質問を生成する** ことです。
質問を生成するには、以下に提供されている [DatabaseSummary] のデータベース要約に頼ることができます。
戦略的な回答のために、**できるだけ多くの関連する質問を生成してください**（最大5つ）。
科学者は **これらの質問を使用して文献を検索する** ので、複雑な質問を少なく作るよりも、シンプルな質問を多く生成する方が望ましいです。

[ DatabaseSummary ]
{ mk_summary }
[/ DatabaseSummary ]
[ HumanQuery ]
{ user_query }
[/ HumanQuery ]

番号付きの質問のみで回答し、このフォーマットに厳密に従ってください。質問の前後に紹介文や結論文は不要です。

例:

1	…
2	…
3	…
…
N …

import torch.nn.functional as F
from sentence_transformers import SentenceTransformer

QUERY_TEMPLATE = "クエリ: {query}"
PASSAGE_TEMPLATE = "文章: {passage}"

model_name = "cl-nagoya/ruri-large"
model = SentenceTransformer(model_name)

sentences = [
    "クエリ: 瑠璃色はどんな色？",
    "文章: 瑠璃色（るりいろ）は、紫みを帯びた濃い青。名は、半貴石の瑠璃（ラピスラズリ、英: lapis lazuli）による。JIS慣用色名では「こい紫みの青」（略号 dp-pB）と定義している[1][2]。",
    "クエリ: ワシやタカのように、鋭いくちばしと爪を持った大型の鳥類を総称して「何類」というでしょう?",
    "文章: ワシ、タカ、ハゲワシ、ハヤブサ、コンドル、フクロウが代表的である。これらの猛禽類はリンネ前後の時代(17~18世紀)には鷲類・鷹類・隼類及び梟類に分類された。ちなみにリンネは狩りをする鳥を単一の目(もく)にまとめ、vultur(コンドル、ハゲワシ)、falco(ワシ、タカ、ハヤブサなど)、strix(フクロウ)、lanius(モズ)の4属を含めている。",
]
embeddings = model.encode(sentences, convert_to_tensor=True)

similarities = F.cosine_similarity(embeddings.unsqueeze(0), embeddings.unsqueeze(1), dim=2)

from sentence_transformers import SentenceTransformer

QUERY_TEMPLATE = "query: {query}"
PASSAGE_TEMPLATE = "passage: {passage}"

model_name = "pkshatech/GLuCoSE-base-ja-v2"
model = SentenceTransformer(model_name)

sentences = [
    'query: PKSHAはどんな会社ですか？'
    'passage: 研究開発したアルゴリズムを、多くの企業のソフトウエア・オペレーションに導入しています。'
]
embeddings = model.encode(sentences)

similarities = F.cosine_similarity(embeddings.unsqueeze(0), embeddings.unsqueeze(1), dim=2)

from sentence_transformers import SentenceTransformer

QUERY_TEMPLATE = "query: {query}"
PASSAGE_TEMPLATE = "passage: {passage}"

model_name = "pkshatech/RoSEtta-base"
model = SentenceTransformer(model_name)

sentences = [
    'query: PKSHAはどんな会社ですか？'
    'passage: 研究開発したアルゴリズムを、多くの企業のソフトウエア・オペレーションに導入しています。'
]
embeddings = model.encode(sentences)

similarities = F.cosine_similarity(embeddings.unsqueeze(0), embeddings.unsqueeze(1), dim=2)

from transformers import AutoModelForMaskedLM,AutoTokenizer
import torch
import numpy as np

model = AutoModelForMaskedLM.from_pretrained("aken12/splade-japanese-v3") 
tokenizer = AutoTokenizer.from_pretrained("aken12/splade-japanese-v3")
vocab_dict = {v: k for k, v in tokenizer.get_vocab().items()}

def encode_query(query): ##query passsage maxlen: 32,180
    query = tokenizer(query, return_tensors="pt")
    output = model(**query, return_dict=True).logits
    output, _ = torch.max(torch.log(1 + torch.relu(output)) * query['attention_mask'].unsqueeze(-1), dim=1)
    return output

with torch.no_grad():
    model_output = encode_query(query="筑波大学では何の研究が行われているか？")

reps = model_output
idx = torch.nonzero(reps[0], as_tuple=False)

dict_splade = {}
for i in idx:
    token_value = reps[0][i[0]].item()
    if token_value > 0:
        token = vocab_dict[int(i[0])]
        dict_splade[token] = float(token_value)

sorted_dict_splade = sorted(dict_splade.items(), key=lambda item: item[1], reverse=True)
for token, value in sorted_dict_splade:
    print(token, value)