michaelfeil / michaelfeil/infinity
float16 and other optimizations help?
Nobody has claimed this yet.
- Dominant language
- Python
- Stars
- 2.9k
- Forks
- 206
- PR merge metrics
- No merged PRs in 30d
Description
I built the following script based on reviewing yours...thanks BTW. It only implements float16 but I'd love to incorporate other quantizations and/or better transformer or whatever else. Hope what I spend a good day on helps you in your endeavors:
```
import torch
from sentence_transformers import SentenceTransformer, util
import numpy as np
import time
class CustomSentenceTransformer(SentenceTransformer):
def __init__(self, model_name_or_path="sentence-transformers/all-mpnet-base-v2", device='cuda', dtype=torch.float16):
# Call the superclass initializer directly with the device. The superclass should handle device placement.
super().__init__(model_name_or_path, device=device)
if torch.cuda.is_available() and device == 'cuda' and dtype == torch.float16:
self.half()
print("Model converted to float16 precision.")
def encode(self, sentences, batch_size=32, show_progress_bar=False, convert_to_numpy=True, normalize_embeddings=True):
self.eval()
all_embeddings = []
for start_index in range(0, len(sentences), batch_size):
sentences_batch = sentences[start_index:start_index+batch_size]
features = self.tokenize(sentences_batch)
features = util.batch_to_device(features, self.device)
with torch.no_grad():
out_features = self.forward(features)
embeddings = out_features['sentence_embedding']
if normalize_embeddings:
embeddings = torch.nn.functional.normalize(embeddings, p=2, dim=1)
all_embeddings.append(embeddings)
all_embeddings = torch.cat(all_embeddings, 0)
if convert_to_numpy:
all_embeddings = all_embeddings.cpu().numpy()
return all_embeddings
```
By the way, I'm noticing 1/2 compute time but the VRAM usage isn't going down at all. The model I'm using was originally created in float16 so...What I'm "guessing" is that even though it's originally float16, prior to using this script, it runs in float32 because dtype isn't set to float16? Guessing that's why you named your class "patched?"
At any rate, would be curious to know your expert opinion. Here's most of the function I'm using with it...
```
def generate_embeddings(text_chunks):
model_name = r"PATH HERE"
model_kwargs = {'device': 'cuda'}
encode_kwargs = {'normalize_embeddings': True}
model = CustomSentenceTransformer(model_name, **model_kwargs)
start_time = time.time()
embeddings = model.encode(text_chunks, **encode_kwargs)
end_time = time.time()
elapsed_time = end_time - start_time
print(f"Embedding generation took {elapsed_time:.2f} seconds.")
```
Contributor guide
No contributing guide indexed for this repository
First steps
- Read the whole issue, then the project's contributing guide.
- Comment on the issue to say you are picking it up — it saves two people doing the same work.
- Fork the repository and make your change on a branch.
- Open a pull request that references the issue number.
Research direction
Start by reviewing the CustomSentenceTransformer class and the generate_embeddings function shown in the issue. Compare the float16 behavior and VRAM usage, then investigate which additional quantization or transformer optimizations are appropriate. Done would require a clearly scoped optimization with measured compute time and memory results.
Written by the indexing model from the issue text.
Assessment
- Tech stack
- python, pytorch
- Domain
- machine-learning
- Issue type
- Feature
- Difficulty
- 5/5
- Estimated time
- Over a week
- Activity status
- Stale
- Clarity
- Needs clarification
- Newbie friendliness
- 20/100