Spaces:
Running
on
T4
Running
on
T4
| from transformers import AutoTokenizer, AutoModel | |
| import torch | |
| import torch.nn.functional as F | |
| import numpy as np | |
| import streamlit as st | |
| import boto3 | |
| import json | |
| runtime = boto3.client('sagemaker-runtime',aws_access_key_id=st.secrets['user_access_key'], | |
| aws_secret_access_key=st.secrets['user_secret_key'],region_name='us-east-1') | |
| # Load Tokenizer from HuggingFace Hub | |
| tokenizer = AutoTokenizer.from_pretrained('sentence-transformers/all-MiniLM-L6-v2') | |
| endpoint_name = 'huggingface-pytorch-inference-2025-05-21-16-31-07-967' | |
| def mean_pooling(token_embeddings, attention_mask): | |
| input_mask_expanded = attention_mask.unsqueeze(-1).expand(token_embeddings.size()).float() | |
| return torch.sum(token_embeddings * input_mask_expanded, dim=1) / \ | |
| torch.clamp(input_mask_expanded.sum(dim=1), min=1e-9) | |
| def vectorise(sentence,token_level_vectors): | |
| encoded_input = tokenizer(sentence, padding=True, truncation=True, return_tensors='pt') | |
| # Get input IDs (token IDs) | |
| input_ids = encoded_input['input_ids'][0] | |
| # Convert IDs to tokens | |
| tokens = tokenizer.convert_ids_to_tokens(input_ids) | |
| model_output = runtime.invoke_endpoint( | |
| EndpointName=endpoint_name, | |
| ContentType="application/json", | |
| Body=json.dumps({"inputs": sentence}) | |
| ) | |
| token_vectors = json.loads(model_output['Body'].read().decode()) | |
| if(token_level_vectors): | |
| return tokens,token_vectors | |
| token_vectors_tensor = torch.tensor(token_vectors) | |
| attention_mask = encoded_input['attention_mask'] | |
| # Perform pooling | |
| sentence_embeddings = mean_pooling(token_vectors_tensor, attention_mask) | |
| # Normalize embeddings | |
| sentence_embeddings = F.normalize(sentence_embeddings, p=2, dim=1) | |
| return sentence_embeddings[0].tolist() | |
| def search(hits): | |
| tokens,token_vectors = vectorise(st.session_state.input_text,True) | |
| final_docs = [] | |
| for ind,j in enumerate(hits): | |
| max_score_dict_list = [] | |
| doc={"_source": | |
| { | |
| "product_description":j["_source"]["description"],"caption":j["_source"]["title"], | |
| "image_url":j["_source"]["image_s3_url"].replace("https://retail-demo-store-us-east-1.s3.amazonaws.com/images/","/home/user/images_retail/") | |
| ,"price":j["_source"]["price"], | |
| "style":j["_source"]["style"],"category":j["_source"]["category"]},"_id":j["_id"],"_score":j["_score"]} | |
| if("gender_affinity" in j["_source"]): | |
| doc["_source"]["gender_affinity"] = j["_source"]["gender_affinity"] | |
| else: | |
| doc["_source"]["gender_affinity"] = "" | |
| source_doc_token_keys = list(j["_source"].keys()) | |
| with_s = [x for x in source_doc_token_keys if x.startswith("description-token-")] | |
| add_score = 0 | |
| for index,i in enumerate(token_vectors[0]): | |
| token = tokens[index] | |
| if(token!='[SEP]' and token!='[CLS]'): | |
| query_token_vector = np.array(i) | |
| scores = [] | |
| for m in with_s: | |
| m_arr = m.split("-") | |
| if(m_arr[-1]!='[SEP]' and m_arr[-1]!='[CLS]'): | |
| doc_token_vector = np.array(j["_source"][m]) | |
| score = np.dot(query_token_vector,doc_token_vector) | |
| scores.append({"doc_token":m_arr[3],"score":score}) | |
| newlist = sorted(scores, key=lambda d: d['score'], reverse=True) | |
| max_score = newlist[0]['score'] | |
| add_score+=max_score | |
| max_score_dict_list.append(newlist[0]) | |
| max_score_dict_list_sorted = sorted(max_score_dict_list, key=lambda d: d['score'], reverse=True) | |
| print(max_score_dict_list_sorted) | |
| doc["total_score"] = add_score | |
| doc['max_score_dict_list_sorted'] = max_score_dict_list_sorted | |
| final_docs.append(doc) | |
| final_docs_sorted = sorted(final_docs, key=lambda d: d['total_score'], reverse=True) | |
| return final_docs_sorted | |