Spaces:
Runtime error
Runtime error
File size: 2,113 Bytes
3539c75 ae1fd26 3539c75 7faa7d5 3539c75 7faa7d5 3539c75 7faa7d5 3539c75 7faa7d5 3539c75 7faa7d5 3539c75 ae1fd26 7faa7d5 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 |
import os
from tqdm import tqdm
from datasets import load_dataset
from sentence_transformers import SentenceTransformer
from qdrant_client import QdrantClient, models
MAX_QUESTIONS = 1000
def compute_embedding(sentences, emb_model):
return emb_model.encode(sentences=sentences)
def get_questions(ds):
questions_text = set()
for i, item in enumerate(ds):
if i == MAX_QUESTIONS:
break
for q_text in item['questions']['text']:
questions_text.add(q_text)
unique_questions = list(questions_text)
return [{'question': q} for q in unique_questions]
def build_index():
qdrant = QdrantClient(
url=os.environ['QDRANT_URL'],
api_key=os.environ['QDRANT_API_KEY'],
)
encoder = SentenceTransformer(model_name_or_path='BAAI/bge-small-en-v1.5')
quora_ds = load_dataset(path='quora', split='train', streaming=True)
quora_questions = get_questions(ds=quora_ds)
qdrant.recreate_collection(
collection_name='questions',
vectors_config=models.VectorParams(
size=encoder.get_sentence_embedding_dimension(),
distance=models.Distance.COSINE
)
)
BATCH_SIZE = 100
question_batch = []
for idx, entry in enumerate(tqdm(quora_questions, desc='Uploading vector embeddings in batch size of {}'.format(BATCH_SIZE))):
if len(question_batch) < BATCH_SIZE:
question_batch.append({
'payload': entry,
'id': idx
})
else:
questions_list = [item['payload']['question'] for item in question_batch]
embedding_batch = compute_embedding(questions_list, encoder).tolist()
records = [
models.Record(
id=entry['id'],
payload=entry['payload'],
vector=embedding
) for entry, embedding in zip(question_batch, embedding_batch)
]
qdrant.upload_records(
collection_name='questions',
records=records
)
question_batch = [] |