face optimizations

This commit is contained in:
mertalev 2023-11-21 01:52:30 -05:00
parent 5cb3312589
commit c1175be1d0
No known key found for this signature in database
GPG key ID: 9181CD92C0A1C5E3
3 changed files with 13 additions and 4 deletions

View file

@ -14,7 +14,7 @@ export class AssetFaceEntity {
@Column({ nullable: true, type: 'uuid' })
personId!: string | null;
@Column({ type: 'float4', array: true })
@Column({ type: 'float4', array: true, select: false })
embedding!: number[];
@Column({ default: 0, type: 'int' })

View file

@ -13,6 +13,7 @@ export class SmartSearchEntity {
@Column({
type: 'float4',
array: true,
select: false,
})
embedding!: number[];
}

View file

@ -15,11 +15,17 @@ import { asVector } from '../infra.utils';
import { dataSource } from '..';
export class PersonRepository implements IPersonRepository {
private readonly faceColumns: string[];
constructor(
@InjectRepository(AssetEntity) private assetRepository: Repository<AssetEntity>,
@InjectRepository(PersonEntity) private personRepository: Repository<PersonEntity>,
@InjectRepository(AssetFaceEntity) private assetFaceRepository: Repository<AssetFaceEntity>,
) {}
) {
this.faceColumns = this.assetFaceRepository.manager.connection
.getMetadata(AssetFaceEntity).ownColumns
.map((column) => column.propertyName)
.filter((propertyName) => propertyName !== 'embedding');
}
/**
* Before reassigning faces, delete potential key violations
@ -245,12 +251,14 @@ export class PersonRepository implements IPersonRepository {
async searchByEmbedding({ ownerId, embedding, numResults, maxDistance }: EmbeddingSearch): Promise<AssetFaceEntity[]> {
const cte = this.assetFaceRepository.createQueryBuilder('faces')
.select('1 + (faces.embedding <=> :embedding)', 'distance')
.leftJoinAndSelect('faces.asset', 'asset')
.innerJoin('faces.asset', 'asset')
.where('asset.ownerId = :ownerId')
.orderBy(`faces.embedding <=> :embedding`)
.setParameters({ownerId, embedding: asVector(embedding)})
.limit(numResults);
this.faceColumns.forEach((col) => cte.addSelect(`faces.${col} AS "${col}"`));
const res = await dataSource.createQueryBuilder()
.select('res.*')
.addCommonTableExpression(cte, 'cte')