This commit is contained in:
Antonio Scandurra 2023-09-07 15:25:23 +02:00
parent 757a285852
commit a45c8c380f

View file

@ -108,21 +108,21 @@ impl EmbeddingQueue {
let finished_files_tx = self.finished_files_tx.clone(); let finished_files_tx = self.finished_files_tx.clone();
let embedding_provider = self.embedding_provider.clone(); let embedding_provider = self.embedding_provider.clone();
self.executor.spawn(async move { self.executor
.spawn(async move {
let mut spans = Vec::new(); let mut spans = Vec::new();
for fragment in &batch { for fragment in &batch {
let file = fragment.file.lock(); let file = fragment.file.lock();
spans.extend( spans.extend(
{
file.spans[fragment.span_range.clone()] file.spans[fragment.span_range.clone()]
.iter().filter(|d| d.embedding.is_none()) .iter()
.map(|d| d.content.clone()) .filter(|d| d.embedding.is_none())
} .map(|d| d.content.clone()),
); );
} }
// If spans is 0, just send the fragment to the finished files if its the last one. // If spans is 0, just send the fragment to the finished files if its the last one.
if spans.len() == 0 { if spans.is_empty() {
for fragment in batch.clone() { for fragment in batch.clone() {
if let Some(file) = Arc::into_inner(fragment.file) { if let Some(file) = Arc::into_inner(fragment.file) {
finished_files_tx.try_send(file.into_inner()).unwrap(); finished_files_tx.try_send(file.into_inner()).unwrap();
@ -135,13 +135,14 @@ impl EmbeddingQueue {
Ok(embeddings) => { Ok(embeddings) => {
let mut embeddings = embeddings.into_iter(); let mut embeddings = embeddings.into_iter();
for fragment in batch { for fragment in batch {
for span in for span in &mut fragment.file.lock().spans[fragment.span_range.clone()]
&mut fragment.file.lock().spans[fragment.span_range.clone()].iter_mut().filter(|d| d.embedding.is_none()) .iter_mut()
.filter(|d| d.embedding.is_none())
{ {
if let Some(embedding) = embeddings.next() { if let Some(embedding) = embeddings.next() {
span.embedding = Some(embedding); span.embedding = Some(embedding);
} else { } else {
log::error!("number of embeddings returned different from number of documents"); log::error!("number of embeddings != number of documents");
} }
} }