Upload folder using huggingface_hub

499de81 verified over 1 year ago

2.02 kB

from typing import Dict, List, Any
import torch
from transformers import AutoTokenizer, AutoModelForSequenceClassification
import torch.nn.functional as F

class EndpointHandler:
    def __init__(self, path: str = "BAAI/bge-reranker-v2-m3"):
        # Load tokenizer and model
        self.tokenizer = AutoTokenizer.from_pretrained(path)
        self.model = AutoModelForSequenceClassification.from_pretrained(path)
        self.model.eval()

        # Determine the computation device
        self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
        self.model.to(self.device)

    def __call__(self, data: Dict[str, Any]) -> List[Dict[str, Any]]:
        """
        Expected input format:
        {
            "query": "Your query here",
            "texts": ["Document 1", "Document 2", ...],
            "normalize": true  # Optional; defaults to False
        }
        """
        query = data.get("query")
        texts = data.get("texts", [])
        normalize = data.get("normalize", False)

        if not query or not texts:
            return [{"error": "Both 'query' and 'texts' fields are required."}]

        # Prepare input pairs
        pairs = [[query, text] for text in texts]

        # Tokenize input pairs
        inputs = self.tokenizer(
            pairs,
            padding=True,
            truncation=True,
            return_tensors="pt",
            max_length=512
        ).to(self.device)

        with torch.no_grad():
            # Get model logits
            outputs = self.model(**inputs)
            scores = outputs.logits.view(-1)

            # Apply sigmoid normalization if requested
            if normalize:
                scores = torch.sigmoid(scores)

        # Prepare the response
        results = [
            {"index": idx, "score": score.item()}
            for idx, score in enumerate(scores)
        ]

        # Sort results by descending score
        results.sort(key=lambda x: x["score"], reverse=True)
        return results

	from typing import Dict, List, Any
	import torch
	from transformers import AutoTokenizer, AutoModelForSequenceClassification
	import torch.nn.functional as F

	class EndpointHandler:
	def __init__(self, path: str = "BAAI/bge-reranker-v2-m3"):
	# Load tokenizer and model
	self.tokenizer = AutoTokenizer.from_pretrained(path)
	self.model = AutoModelForSequenceClassification.from_pretrained(path)
	self.model.eval()

	# Determine the computation device
	self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
	self.model.to(self.device)

	def __call__(self, data: Dict[str, Any]) -> List[Dict[str, Any]]:
	"""
	Expected input format:
	{
	"query": "Your query here",
	"texts": ["Document 1", "Document 2", ...],
	"normalize": true # Optional; defaults to False
	}
	"""
	query = data.get("query")
	texts = data.get("texts", [])
	normalize = data.get("normalize", False)

	if not query or not texts:
	return [{"error": "Both 'query' and 'texts' fields are required."}]

	# Prepare input pairs
	pairs = [[query, text] for text in texts]

	# Tokenize input pairs
	inputs = self.tokenizer(
	pairs,
	padding=True,
	truncation=True,
	return_tensors="pt",
	max_length=512
	).to(self.device)

	with torch.no_grad():
	# Get model logits
	outputs = self.model(**inputs)
	scores = outputs.logits.view(-1)

	# Apply sigmoid normalization if requested
	if normalize:
	scores = torch.sigmoid(scores)

	# Prepare the response
	results = [
	{"index": idx, "score": score.item()}
	for idx, score in enumerate(scores)
	]

	# Sort results by descending score
	results.sort(key=lambda x: x["score"], reverse=True)
	return results