deploy

service/prediction_service.py

@@ -19,6 +19,11 @@ if not os.path.exists(os.path.join(model_dir, 'config.json')):
     model.save(model_dir)
 else:
     print(f"Loading model from local directory: {model_dir}")
-    # model = SentenceTransformer(model_dir)
-    model = SentenceTransformer(model_dir, trust_remote_code=True)
+    model = SentenceTransformer(model_dir, trust_remote_code=True)  # Added trust_remote_code=True
 
+# Define predict_label function
+def predict_label(text):
+    # Implement your prediction logic using the model
+    embeddings = model.encode([text])
+    # For demonstration, let's return the embedding shape
+    return embeddings.shape