aakarsh-yadav-tcgls commited on
Commit
fe3da63
·
1 Parent(s): 5a6d285

update ignore file

Browse files
Files changed (1) hide show
  1. app_api.py +4 -1
app_api.py CHANGED
@@ -12,6 +12,9 @@ model_name = "openlm-research/open_llama_3b_v2"
12
  tokenizer = LlamaTokenizer.from_pretrained(model_name)
13
  model = LlamaModel.from_pretrained(model_name, from_safetensors=False)
14
  # model.eval()
 
 
 
15
 
16
  app = Flask(__name__)
17
 
@@ -23,7 +26,7 @@ def generate_embeddings():
23
 
24
  embeddings = []
25
  for product in products:
26
- inputs = tokenizer(product["description"], return_tensors="pt", truncation=True, max_length=512)
27
  with torch.no_grad():
28
  outputs = model(**inputs, output_hidden_states=True)
29
  embedding = outputs.hidden_states[-1].mean(dim=1).squeeze().numpy()
 
12
  tokenizer = LlamaTokenizer.from_pretrained(model_name)
13
  model = LlamaModel.from_pretrained(model_name, from_safetensors=False)
14
  # model.eval()
15
+ device = "cuda" if torch.cuda.is_available() else "cpu"
16
+
17
+ model = model.to(device)
18
 
19
  app = Flask(__name__)
20
 
 
26
 
27
  embeddings = []
28
  for product in products:
29
+ inputs = tokenizer(product["description"], return_tensors="pt", truncation=True, max_length=512).to(device)
30
  with torch.no_grad():
31
  outputs = model(**inputs, output_hidden_states=True)
32
  embedding = outputs.hidden_states[-1].mean(dim=1).squeeze().numpy()