Spaces:
Paused
Paused
aakarsh-yadav-tcgls commited on
Commit ·
fe3da63
1
Parent(s): 5a6d285
update ignore file
Browse files- 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()
|