Update app.py

app.py

@@ -39,6 +39,17 @@ model_name = "daekeun-ml/Llama-2-ko-instruct-13B"
 tokenizer = AutoTokenizer.from_pretrained(model_name)
 model = AutoDistributedModelForCausalLM.from_pretrained(model_name)
 
+
+def check(model_name):
+    data = requests.get("https://health.petals.dev/api/v1/state").json()
+    out = []
+    for d in data['model_reports']:
+        if d['name'] == model_name:
+            if d['state']=="healthy":
+                return True
+    return False
+
+
 def chat(id, npc, prompt):
 
     # get_coin endpoint
@@ -52,13 +63,13 @@ def chat(id, npc, prompt):
         return "no coin"
     
     # model inference
-
-    prom = ""
-    inputs = tokenizer(prom, return_tensors="pt")["input_ids"]
-    outputs = model.generate(inputs, max_new_tokens=100)
-    print(tokenizer.decode(outputs[0]))
-    
-    output = "AI 응답입니다."
+    if check:
+        prom = ""
+        inputs = tokenizer(prom, return_tensors="pt")["input_ids"]
+        outputs = model.generate(inputs, max_new_tokens=100)
+        print(tokenizer.decode(outputs[0]))
+    else:
+        output = "no model"
 
 
     # add_transaction endpoint