Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

feat(cli serve): add a new reranker BCEReranking #25

Merged
merged 1 commit into from
Mar 1, 2024

Conversation

bjwswang
Copy link
Collaborator

@bjwswang bjwswang commented Mar 1, 2024

What type of PR is this?

What this PR does / why we need it

Which issue(s) this PR fixes

Fixes #

Special notes for your reviewer

if isinstance(result, float):
result = [result]
return result
else:
return None


def select_reranking(model_name_or_path: str) -> BaseReranking:

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

look like if we use transformers, we can have same code to serve diffrent rerank model?

import torch
from transformers import AutoTokenizer, AutoModelForSequenceClassification

# init model and tokenizer
tokenizer = AutoTokenizer.from_pretrained('maidalun1020/bce-reranker-base_v1')
model = AutoModelForSequenceClassification.from_pretrained('maidalun1020/bce-reranker-base_v1')

device = 'cuda'  # if no GPU, set "cpu"
model.to(device)

# get inputs
inputs = tokenizer(sentence_pairs, padding=True, truncation=True, max_length=512, return_tensors="pt")
inputs_on_device = {k: v.to(device) for k, v in inputs.items()}

# calculate scores
scores = model(**inputs_on_device, return_dict=True).logits.view(-1,).float()
scores = torch.sigmoid(scores)
import torch
from transformers import AutoModelForSequenceClassification, AutoTokenizer

tokenizer = AutoTokenizer.from_pretrained('BAAI/bge-reranker-large')
model = AutoModelForSequenceClassification.from_pretrained('BAAI/bge-reranker-large')
model.eval()

pairs = [['what is panda?', 'hi'], ['what is panda?', 'The giant panda (Ailuropoda melanoleuca), sometimes called a panda bear or simply panda, is a bear species endemic to China.']]
with torch.no_grad():
    inputs = tokenizer(pairs, padding=True, truncation=True, return_tensors='pt', max_length=512)
    scores = model(**inputs, return_dict=True).logits.view(-1, ).float()
    print(scores)

Copy link
Collaborator Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

FlagEmbedding/BCEEmbedding provides interfaces to load reranking models which uses the transformers as well. By using FlagEmbedding/BCEEmbedding,we do not need to care much about the base dependecies.

@bjwswang bjwswang merged commit c164466 into kubeagi:main Mar 1, 2024
2 of 3 checks passed
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Labels
None yet
Projects
None yet
Development

Successfully merging this pull request may close these issues.

2 participants