chryzxc commited on
Commit
6d7e0c5
·
verified ·
1 Parent(s): dd92c0c

Update app.py

Browse files
Files changed (1) hide show
  1. app.py +17 -11
app.py CHANGED
@@ -39,21 +39,27 @@ async def predict(request: Request):
39
  raise HTTPException(status_code=400, detail="No text provided")
40
 
41
  # Tokenize input
42
- inputs = tokenizer(text)
 
 
 
 
 
 
43
 
44
- # Run model
45
- outputs = session.run(None, {
46
- "input_ids": inputs["input_ids"].astype(np.int64),
47
- "attention_mask": inputs["attention_mask"].astype(np.int64)
48
- })
 
49
 
50
  # Prepare response with converted types
51
- response = {
52
- "embedding": convert_output(outputs[0]), # Process main output
53
- "tokens": tokenizer.convert_ids_to_tokens(inputs["input_ids"][0])
 
54
  }
55
- print("embeddings", response["embedding"])
56
- return jsonable_encoder(response)
57
 
58
  except Exception as e:
59
  raise HTTPException(status_code=500, detail=str(e))
 
39
  raise HTTPException(status_code=400, detail="No text provided")
40
 
41
  # Tokenize input
42
+ inputs = tokenizer(
43
+ text,
44
+ return_tensors="np",
45
+ padding=False, # Disable padding
46
+ truncation=False, # Disable truncation
47
+ add_special_tokens=True # Ensure CLS/SEP tokens
48
+ )
49
 
50
+ onnx_inputs = {
51
+ "input_ids": np.array(inputs["input_ids"], dtype=np.int64),
52
+ "attention_mask": np.array(inputs["attention_mask"], dtype=np.int64)
53
+ }
54
+
55
+ outputs = session.run(None, onnx_inputs)
56
 
57
  # Prepare response with converted types
58
+ return {
59
+ "embedding": outputs[0][0].astype(float).tolist(),
60
+ "input_ids": inputs["input_ids"][0].tolist(),
61
+ "attention_mask": inputs["attention_mask"][0].tolist()
62
  }
 
 
63
 
64
  except Exception as e:
65
  raise HTTPException(status_code=500, detail=str(e))