Skip to content

Commit

Permalink
Merge pull request #296 from Samagra-Development/rerank
Browse files Browse the repository at this point in the history
added bge reranker
  • Loading branch information
Gautam-Rajeev authored Feb 28, 2024
2 parents 7a21229 + 260025a commit 695419d
Show file tree
Hide file tree
Showing 7 changed files with 104 additions and 0 deletions.
15 changes: 15 additions & 0 deletions src/rerankers/bge_base/local/Dockerfile
Original file line number Diff line number Diff line change
@@ -0,0 +1,15 @@
# Use an official Python runtime as a parent image
FROM python:3.9-slim

WORKDIR /app


#install requirements
COPY requirements.txt requirements.txt
RUN pip3 install -r requirements.txt

# Copy the rest of the application code to the working directory
COPY . /app/
EXPOSE 8000
# Set the entrypoint for the container
CMD ["hypercorn", "--bind", "0.0.0.0:8000", "api:app"]
20 changes: 20 additions & 0 deletions src/rerankers/bge_base/local/README.md
Original file line number Diff line number Diff line change
@@ -0,0 +1,20 @@
## Grievance classification:


### Purpose :
Re rank given a question and a list of contetn


### Testing the model deployment :
To run for testing just the Hugging Face deployment for grievence recognition, you can follow the following steps :

- Git clone the repo
- Go to current folder location i.e. ``` cd /src/rerankers/bge_base/local ```
- Create docker image file and test the api:
```
docker build -t testmodel .
docker run -p 8000:8000 testmodel
curl -X POST -H "Content-Type: application/json" \
-d '{"question": "What is agriculture ?", "content_chunks": ["Farming is a practice of growing crops to sell them to generate money", "LLM are the present day hype machine but will they be useful until you can truly reason with them? ", "Things are generally better than what people deep into it feel"]}' \
http://localhost:8000/
```
2 changes: 2 additions & 0 deletions src/rerankers/bge_base/local/__init__.py
Original file line number Diff line number Diff line change
@@ -0,0 +1,2 @@
from .request import ModelRequest
from .request import Model
30 changes: 30 additions & 0 deletions src/rerankers/bge_base/local/api.py
Original file line number Diff line number Diff line change
@@ -0,0 +1,30 @@
from model import Model
from request import ModelRequest
from quart import Quart, request,jsonify
import aiohttp

app = Quart(__name__)

model = None

@app.before_serving
async def startup():
app.client = aiohttp.ClientSession()
global model
model = Model(app)

@app.route('/', methods=['POST'])
async def embed():
global model
data = await request.get_json()
req = ModelRequest(**data)
prediction = await model.inference(req)
# Convert the NumPy array to a list (or another serializable format) and return as JSON
if prediction is not None:
return jsonify(prediction.tolist()) # Assuming 'prediction' is a NumPy array
else:
# Return a meaningful error message if prediction is None
return jsonify({'error': 'Prediction failed'}), 500

if __name__ == "__main__":
app.run()
18 changes: 18 additions & 0 deletions src/rerankers/bge_base/local/model.py
Original file line number Diff line number Diff line change
@@ -0,0 +1,18 @@
from request import ModelRequest
from sentence_transformers.cross_encoder import CrossEncoder
import torch

class Model():
def __new__(cls, context):
cls.context = context
if not hasattr(cls, 'instance'):
cls.instance = super(Model, cls).__new__(cls)
model_name = "BAAI/bge-reranker-base"
cls.model = CrossEncoder(model_name)
return cls.instance


async def inference(self, request: ModelRequest):
predict_array = request.predict_array
predictions = self.model.predict(predict_array)
return (predictions)
15 changes: 15 additions & 0 deletions src/rerankers/bge_base/local/request.py
Original file line number Diff line number Diff line change
@@ -0,0 +1,15 @@
import requests
import json

class ModelRequest():
def __init__(self, question, content_chunks):
self.question = question
self.content_chunks = content_chunks
self.predict_array = [[question, content] for content in content_chunks]

def to_json(self):
return json.dumps(self, default=lambda o: o.__dict__,
sort_keys=True, indent=4)



4 changes: 4 additions & 0 deletions src/rerankers/bge_base/local/requirements.txt
Original file line number Diff line number Diff line change
@@ -0,0 +1,4 @@
torch
sentence_transformers
quart
aiohttp

0 comments on commit 695419d

Please sign in to comment.