update app logic

app.py

@@ -1,12 +1,11 @@
 from flask import Flask, request, jsonify
-from transformers import AutoModel, AutoTokenizer
-import torch
+from transformers import TFAutoModel, AutoTokenizer
 
 app = Flask(__name__)
 
-# Load PhoBERT
+# Load PhoBERT (TensorFlow version)
 tokenizer = AutoTokenizer.from_pretrained("vinai/phobert-base")
-model = AutoModel.from_pretrained("vinai/phobert-base", from_tf=True)
+model = TFAutoModel.from_pretrained("vinai/phobert-base", from_tf=True)
 
 @app.route('/embed', methods=['POST'])
 def embed():
@@ -15,12 +14,11 @@ def embed():
     if not text:
         return jsonify({"error": "No text provided"}), 400
 
-    inputs = tokenizer(text, return_tensors="pt")
-    with torch.no_grad():
-        outputs = model(**inputs)
+    inputs = tokenizer(text, return_tensors="tf")  # Chuyển sang TensorFlow tensor
+    outputs = model(**inputs)
 
     # Lấy embedding từ hidden state đầu tiên
-    embedding = outputs.last_hidden_state[:, 0, :].squeeze().tolist()
+    embedding = outputs.last_hidden_state[:, 0, :].numpy().tolist()  # Dùng .numpy() để chuyển từ TensorFlow tensor sang list
 
     return jsonify({"embedding": embedding})