fix

app.py

@@ -9,16 +9,20 @@ app = FastAPI()
 class Input(BaseModel):
     text1 : List
     text2 : List
+    model : str = "sentence-transformers/all-MiniLM-L6-v2"
 
 class Output(BaseModel):
     matrix : List
 
-@app.post("/process", response_model=Output)
+@app.post("/saoke-to-heatmap", response_model=Output)
 def saoke_to_heatmap(payload: Input):
     saoke_spec = text_to_saoke(payload.text1)
     saoke_patent = text_to_saoke(payload.text2)
 
-    model = SentenceTransformer('sentence-transformers/all-MiniLM-L6-v2')
+    try:
+        model = SentenceTransformer(payload.model)
+    except:
+        model = SentenceTransformer('sentence-transformers/all-MiniLM-L6-v2')
 
     embeddings1, embeddings2 = text_to_embeddings(saoke_spec, saoke_patent, model)
     
@@ -28,7 +32,12 @@ def saoke_to_heatmap(payload: Input):
 
 @app.post("/text-to-heatmap", response_model=Output)
 def text_to_heatmap(payload: Input):
-    embeddings1, embeddings2 = text_to_embeddings(payload.text1, payload.text2)
+    try:
+        model = SentenceTransformer(payload.model)
+    except:
+        model = SentenceTransformer('sentence-transformers/all-MiniLM-L6-v2')
+
+    embeddings1, embeddings2 = text_to_embeddings(payload.text1, payload.text2, model)
     matrix = embeddings_to_matrix(embeddings1, embeddings2)
     print({"matrix": matrix})
     return {"matrix": matrix}