Update main.py

main.py

@@ -142,6 +142,11 @@ def getChat(text,model,tokenizer):
             size_list.append((start,end))
             start = end
     # print(size_list)
+    size_str = ""
+    if len(size_list)==0:
+        size_str = "无"
+    else:
+        size_str = ''.join(size_list)
 
     cancer_dict = {'腺癌': 0, '肺良性疾病': 1, '鳞癌': 2, '无法判断组织分型': 3, '复合型': 4, '转移癌': 5, '小细胞癌': 6, '大细胞癌': 7}
     id_cancer = {j:i for i,j in cancer_dict.items()}
@@ -153,9 +158,10 @@ def getChat(text,model,tokenizer):
     cancer = id_cancer[cancer[0]]
     transfer = id_transfer[transfer[0]]
     ly_transfer = id_lymph_transfer[ly_transfer[0]]
-    print(cancer,transfer,ly_transfer)
+    # print(cancer,transfer,ly_transfer)
+    output = "肿瘤大小:"+size_str+" <4cm"+"\n"+"病理组织分型"+cancer+"\n"+"转移"+transfer+"\n"+"淋巴转移"+ly_transfer+"\n"
 
-    return size_list,cancer,transfer,ly_transfer
+    return output
 
 app = FastAPI()
 
@@ -177,9 +183,9 @@ tokenizer,model_extra = model_init()
 
 @app.get("/infer_t5")
 def t5(input):
-    r1,r2,r3,r4 = getChat(input,model_extra,tokenizer)
+    output = getChat(input,model_extra,tokenizer)
     # output = pipe_flan(input)
-    return {"output": r4}
+    return {"output": output}
 
 app.mount("/", StaticFiles(directory="static", html=True), name="static")