Update app.py

app.py

@@ -1,13 +1,13 @@
 ##########====================================================================################
 ##########====================PRODUCTION VERSION -- vLLM, GRADIO=====================###########
 ##########====================================================================################
-import os
+import os, subprocess, threading, time
 import requests
 from typing import List, Dict, Tuple
 from datetime import datetime
 from anthropic import Anthropic
 from openai import OpenAI
-import time, gradio as gr
+import gradio as gr
 from tqdm import tqdm
 
 ANTHROPIC_API_KEY = os.getenv("ANTHROPIC_API_KEY")
@@ -22,6 +22,18 @@ open_source_client = OpenAI(api_key="EMPTY", base_url=VLLM_API)
 claude_client = Anthropic(api_key=ANTHROPIC_API_KEY)
 
 
+def wait_for_vllm_ready(timeout=120):
+    start = time.time() 
+    while time.time() - start < timeout:
+        try:
+            r = requests.get("http://localhost:8000/health", timeout=3)
+            if r.status_code == 200:
+                return True
+        except Exception:
+            pass
+        time.sleep(2)
+    raise RuntimeError("vLLM did not start within timeout")
+
 def invoke_messages(
         rows_num: int,
         business_category: str,
@@ -76,39 +88,24 @@ def pass_claude_msg(file_format: str, content: str) -> Tuple[str, str]:
 
 
 def generate_output(messages):
-    enable_model()
-
-    inputs = tokenizer.apply_chat_template(
-        messages,
-        return_tensors="pt",
-        return_dict=True,  ### IMPORTANT: to get a mapping
-        tokenize=True,
-        add_generation_prompt=True,
-        padding=True,
-        return_attention_mask=True
-    ).to(model.device)
-
-    # print(inputs)
-
-    outputs = model.generate(
-        **inputs,
-        max_new_tokens=400,
-        temperature=0.2
+    
+    resp = open_source_client.chat.completions.create(
+        model=QWEN_MODEL,
+        messages=messages,
+        max_tokens=400,
+        temperature=0.2,
+        stream=False
     )
 
-    ### Get the length(num of tokens) of the input prompt
-    prompt_len = inputs["input_ids"].shape[1]
-
-    ### Slice the generated sequence to skip the prompt length
-    gen_tokens = outputs[0][prompt_len:]
-
-    # print(tokenizer.decode(gen_tokens, skip_special_tokens=True))
-
-    return gen_tokens
+    return resp.choices[0].message.content
 
 
 def launch_claude_api(sys_msg, user_msg):
-    response = claude.messages.create(
+
+    if not claude_client:
+        return None
+    
+    response = claude_client.messages.create(
         model=CLAUDE_MODEL,
         system=sys_msg,
         max_tokens=400,
@@ -124,6 +121,7 @@ def launch_claude_api(sys_msg, user_msg):
 
 def generate_mockup_data(category, num_data_rows, columns, a_instruction,
                          progress=gr.Progress()):
+    wait_for_vllm_ready()
     progress(0.2, desc="Generating...")
     msg = invoke_messages(
         rows_num=int(num_data_rows or 10),
@@ -132,16 +130,10 @@ def generate_mockup_data(category, num_data_rows, columns, a_instruction,
         instruction=a_instruction
     )
 
-    resp = open_source_client.chat.completions.create(
-        model=QWEN_MODEL,
-        messages=msg,
-        max_tokens=400,
-        temperature=0.2,
-        stream=False
-    )
+    result = generate_output(msg)
     progress(1.0, desc="Done")
 
-    return resp.choices[0].message.content
+    return result
 
 
 def show_hidden_row():
@@ -158,6 +150,9 @@ def make_file(btn_sort: str, category: str, content: str):
     if not content or not content.strip():
         raise gr.Error("The result content is empty. Cannot create a file.")
 
+    if not claude_client:
+        raise gr.Error("File formatting requires ANTHROPIC_API_KEY.")
+
     try:
         sys_msg, user_msg = pass_claude_msg(btn_sort, content)
         claude_output = launch_claude_api(sys_msg, user_msg)
@@ -242,12 +237,6 @@ def render_interface():
             outputs=btn_json
         )
 
-        ### Pre-warming the model right upon the page load
-        ### in order to save the model load time when user submitting the form.
-        demo.load(lambda: enable_model(), queue=False)
-
-    return demo
-
 
 if __name__ == "__main__":
     app = render_interface()