Truncate image queries below max tokens length supported by ML model

This would previously return the infamous tensor size mismatch error
Verify this error is not raised since adding the query truncation logic
This commit is contained in:
Debanjum Singh Solanky
2023-01-21 14:11:00 -03:00
parent 3d9ed91e42
commit 6908b6eed3
2 changed files with 26 additions and 1 deletions

View File

@@ -143,7 +143,9 @@ def query(raw_query, count, model: ImageSearchModel):
query.thumbnail((640, query.height)) # scale down image for faster processing
logger.info(f"Find Images by Image: {query_imagepath}")
else:
query = raw_query
# Truncate words in query to stay below max_tokens supported by ML model
max_words = 20
query = " ".join(raw_query.split()[:max_words])
logger.info(f"Find Images by Text: {query}")
# Now we encode the query (which can either be an image or a text string)