Update app.py

app.py

@@ -3,6 +3,7 @@ import os
 from starlette.middleware.cors import CORSMiddleware
 from pydantic import BaseModel
 import uvicorn
+from typing import List
 import torch
 from fastapi.encoders import jsonable_encoder
 from fastapi.responses import JSONResponse
@@ -18,7 +19,7 @@ app.add_middleware(
 )
 
 class Query(BaseModel):
-    query :str
+    smiles :List[str]
 
 rob_chem_model = ClassificationModel('roberta', 'seyonec/SMILES_tokenized_PubChem_shard00_160k',use_cuda=False ,args={'evaluate_each_epoch':True , 'evaluate_during_training_verbose':True, 'seed':4})
 
@@ -26,14 +27,16 @@ rob_chem_model = ClassificationModel('roberta', 'seyonec/SMILES_tokenized_PubChe
 async def c(query:Query):
     print(query)
     print(type(query))
+    query_dict = query.dict()
     try:
-        predictions, raw_outputs = rob_chem_model.predict([str(query.query)])
+        predictions, raw_outputs = rob_chem_model.predict([str(query.smiles)])
         print(predictions)
         #json_compatible_item_data = jsonable_encoder(predictions[0])
         #return JSONResponse(content=json_compatible_item_data)
-        answer = {"prediction":predictions[0]}
-        return {'response':answer}
-        
+        query_dict.update({"predictions":predictions})
+        #answer = {"prediction":predictions[0]}
+        #return {'response':answer}
+        return query_dict
 
     except Exception as e:
         print("Excepted")