Switch to vLLM for faster inference with lazy loading and multi-turn fix

app.py

@@ -1,77 +1,83 @@
-# app.py
+# app_vllm.py - Faster inference using vLLM
 import os
-import torch
 import spaces
 import gradio as gr
-from transformers import AutoTokenizer, AutoModelForCausalLM
+from vllm import LLM, SamplingParams
+from vllm.lora.request import LoRARequest
+from transformers import AutoTokenizer
 
 HF_TOKEN = os.getenv("HF_TOKEN") or os.getenv("HUGGINGFACEHUB_API_TOKEN")
 
 BASE_MODEL_ID = "meta-llama/Meta-Llama-3.1-8B-Instruct"
 PEFT_MODEL_ID = "befm/Be.FM-8B"
 
-USE_PEFT = True
-try:
-    from peft import PeftModel, PeftConfig  # noqa
-except Exception:
-    USE_PEFT = False
-    print("[WARN] 'peft' not installed; running base model only.")
-
-def load_model_and_tokenizer():
+def load_model():
     if HF_TOKEN is None:
         raise RuntimeError(
             "HF_TOKEN is not set. Add it in Space → Settings → Secrets. "
             "Also ensure your account has access to the gated base model."
         )
-    dtype = torch.float16 if torch.cuda.is_available() else torch.float32
-    tok = AutoTokenizer.from_pretrained(BASE_MODEL_ID, token=HF_TOKEN)
-    if tok.pad_token is None:
-        tok.pad_token = tok.eos_token
-
-    base = AutoModelForCausalLM.from_pretrained(
-        BASE_MODEL_ID,
-        device_map="auto" if torch.cuda.is_available() else None,
-        torch_dtype=dtype,
-        token=HF_TOKEN,
+
+    # Initialize vLLM with PEFT support
+    llm = LLM(
+        model=BASE_MODEL_ID,
+        tokenizer=BASE_MODEL_ID,
+        enable_lora=True,
+        max_lora_rank=64,
+        dtype="float16",
+        gpu_memory_utilization=0.9,
+        trust_remote_code=True,
+    )
+
+    print(f"[INFO] vLLM loaded base model: {BASE_MODEL_ID}")
+
+    # Load PEFT adapter
+    lora_request = LoRARequest(
+        lora_name="befm",
+        lora_int_id=1,
+        lora_path=PEFT_MODEL_ID,
     )
+    print(f"[INFO] PEFT adapter prepared: {PEFT_MODEL_ID}")
+
+    return llm, lora_request
 
-    if USE_PEFT:
-        try:
-            _ = PeftConfig.from_pretrained(PEFT_MODEL_ID, token=HF_TOKEN)
-            model = PeftModel.from_pretrained(base, PEFT_MODEL_ID, token=HF_TOKEN)
-            print(f"[INFO] Loaded PEFT adapter: {PEFT_MODEL_ID}")
-            return model, tok
-        except Exception as e:
-            print(f"[WARN] Failed to load PEFT adapter: {e}")
-            return base, tok
-    return base, tok
+# Lazy load model and tokenizer
+_llm = None
+_lora_request = None
+_tokenizer = None
 
-model, tokenizer = load_model_and_tokenizer()
-DEVICE = model.device
+def get_model_and_tokenizer():
+    global _llm, _lora_request, _tokenizer
+    if _llm is None:
+        _llm, _lora_request = load_model()
+        _tokenizer = AutoTokenizer.from_pretrained(BASE_MODEL_ID, token=HF_TOKEN)
+    return _llm, _lora_request, _tokenizer
 
 @spaces.GPU
-@torch.inference_mode()
 def generate_response(messages, max_new_tokens=512, temperature=0.7, top_p=0.9) -> str:
+    llm, lora_request, tokenizer = get_model_and_tokenizer()
+
     # Apply Llama 3.1 chat template
     prompt = tokenizer.apply_chat_template(
         messages,
         tokenize=False,
         add_generation_prompt=True
     )
-    enc = tokenizer(prompt, return_tensors="pt", padding=True, truncation=True)
-    enc = {k: v.to(DEVICE) for k, v in enc.items()}
 
-    input_length = enc['input_ids'].shape[1]
-    out = model.generate(
-        **enc,
-        max_new_tokens=max_new_tokens,
-        do_sample=True,
+    sampling_params = SamplingParams(
         temperature=temperature,
         top_p=top_p,
-        pad_token_id=tokenizer.eos_token_id,
+        max_tokens=max_new_tokens,
     )
-    # Decode only the newly generated tokens
-    return tokenizer.decode(out[0][input_length:], skip_special_tokens=True)
+
+    # Generate with vLLM
+    outputs = llm.generate(
+        prompts=[prompt],
+        sampling_params=sampling_params,
+        lora_request=lora_request,
+    )
+
+    return outputs[0].outputs[0].text
 
 def chat_fn(message, history, system_prompt, max_new_tokens, temperature, top_p):
     # Build conversation in Llama 3.1 chat format
@@ -103,8 +109,8 @@ demo = gr.ChatInterface(
         gr.Slider(0.1, 1.5, value=0.7, step=0.05, label="temperature"),
         gr.Slider(0.1, 1.0, value=0.9, step=0.05, label="top_p"),
     ],
-    title="Be.FM-8B (PEFT) on Meta-Llama-3.1-8B-Instruct",
-    description="Chat interface using Meta-Llama-3.1-8B-Instruct with PEFT adapter befm/Be.FM-8B."
+    title="Be.FM-8B (vLLM) - Fast Inference",
+    description="Chat interface using vLLM for optimized inference with Meta-Llama-3.1-8B-Instruct and PEFT adapter befm/Be.FM-8B."
 )
 
 if __name__ == "__main__":

app_transformers.py

@@ -0,0 +1,111 @@
+# app.py
+import os
+import torch
+import spaces
+import gradio as gr
+from transformers import AutoTokenizer, AutoModelForCausalLM
+
+HF_TOKEN = os.getenv("HF_TOKEN") or os.getenv("HUGGINGFACEHUB_API_TOKEN")
+
+BASE_MODEL_ID = "meta-llama/Meta-Llama-3.1-8B-Instruct"
+PEFT_MODEL_ID = "befm/Be.FM-8B"
+
+USE_PEFT = True
+try:
+    from peft import PeftModel, PeftConfig  # noqa
+except Exception:
+    USE_PEFT = False
+    print("[WARN] 'peft' not installed; running base model only.")
+
+def load_model_and_tokenizer():
+    if HF_TOKEN is None:
+        raise RuntimeError(
+            "HF_TOKEN is not set. Add it in Space → Settings → Secrets. "
+            "Also ensure your account has access to the gated base model."
+        )
+    dtype = torch.float16 if torch.cuda.is_available() else torch.float32
+    tok = AutoTokenizer.from_pretrained(BASE_MODEL_ID, token=HF_TOKEN)
+    if tok.pad_token is None:
+        tok.pad_token = tok.eos_token
+
+    base = AutoModelForCausalLM.from_pretrained(
+        BASE_MODEL_ID,
+        device_map="auto" if torch.cuda.is_available() else None,
+        torch_dtype=dtype,
+        token=HF_TOKEN,
+    )
+
+    if USE_PEFT:
+        try:
+            _ = PeftConfig.from_pretrained(PEFT_MODEL_ID, token=HF_TOKEN)
+            model = PeftModel.from_pretrained(base, PEFT_MODEL_ID, token=HF_TOKEN)
+            print(f"[INFO] Loaded PEFT adapter: {PEFT_MODEL_ID}")
+            return model, tok
+        except Exception as e:
+            print(f"[WARN] Failed to load PEFT adapter: {e}")
+            return base, tok
+    return base, tok
+
+model, tokenizer = load_model_and_tokenizer()
+DEVICE = model.device
+
+@spaces.GPU
+@torch.inference_mode()
+def generate_response(messages, max_new_tokens=512, temperature=0.7, top_p=0.9) -> str:
+    # Apply Llama 3.1 chat template
+    prompt = tokenizer.apply_chat_template(
+        messages,
+        tokenize=False,
+        add_generation_prompt=True
+    )
+    enc = tokenizer(prompt, return_tensors="pt", padding=True, truncation=True)
+    enc = {k: v.to(DEVICE) for k, v in enc.items()}
+
+    input_length = enc['input_ids'].shape[1]
+    out = model.generate(
+        **enc,
+        max_new_tokens=max_new_tokens,
+        do_sample=True,
+        temperature=temperature,
+        top_p=top_p,
+        pad_token_id=tokenizer.eos_token_id,
+    )
+    # Decode only the newly generated tokens
+    return tokenizer.decode(out[0][input_length:], skip_special_tokens=True)
+
+def chat_fn(message, history, system_prompt, max_new_tokens, temperature, top_p):
+    # Build conversation in Llama 3.1 chat format
+    messages = []
+    if system_prompt:
+        messages.append({"role": "system", "content": system_prompt})
+
+    # History is already in dict format: [{"role": "user", "content": "..."}, ...]
+    for msg in (history or []):
+        messages.append(msg)
+
+    if message:
+        messages.append({"role": "user", "content": message})
+
+    reply = generate_response(
+        messages,
+        max_new_tokens=max_new_tokens,
+        temperature=temperature,
+        top_p=top_p,
+    )
+    return reply
+
+demo = gr.ChatInterface(
+    fn=lambda message, history, system_prompt, max_new_tokens, temperature, top_p:
+        chat_fn(message, history, system_prompt, max_new_tokens, temperature, top_p),
+    additional_inputs=[
+        gr.Textbox(label="System prompt (optional)", placeholder="You are Be.FM assistant...", lines=2),
+        gr.Slider(16, 2048, value=512, step=16, label="max_new_tokens"),
+        gr.Slider(0.1, 1.5, value=0.7, step=0.05, label="temperature"),
+        gr.Slider(0.1, 1.0, value=0.9, step=0.05, label="top_p"),
+    ],
+    title="Be.FM-8B (PEFT) on Meta-Llama-3.1-8B-Instruct",
+    description="Chat interface using Meta-Llama-3.1-8B-Instruct with PEFT adapter befm/Be.FM-8B."
+)
+
+if __name__ == "__main__":
+    demo.launch()

app_vllm.py

@@ -0,0 +1,117 @@
+# app_vllm.py - Faster inference using vLLM
+import os
+import spaces
+import gradio as gr
+from vllm import LLM, SamplingParams
+from vllm.lora.request import LoRARequest
+from transformers import AutoTokenizer
+
+HF_TOKEN = os.getenv("HF_TOKEN") or os.getenv("HUGGINGFACEHUB_API_TOKEN")
+
+BASE_MODEL_ID = "meta-llama/Meta-Llama-3.1-8B-Instruct"
+PEFT_MODEL_ID = "befm/Be.FM-8B"
+
+def load_model():
+    if HF_TOKEN is None:
+        raise RuntimeError(
+            "HF_TOKEN is not set. Add it in Space → Settings → Secrets. "
+            "Also ensure your account has access to the gated base model."
+        )
+
+    # Initialize vLLM with PEFT support
+    llm = LLM(
+        model=BASE_MODEL_ID,
+        tokenizer=BASE_MODEL_ID,
+        enable_lora=True,
+        max_lora_rank=64,
+        dtype="float16",
+        gpu_memory_utilization=0.9,
+        trust_remote_code=True,
+    )
+
+    print(f"[INFO] vLLM loaded base model: {BASE_MODEL_ID}")
+
+    # Load PEFT adapter
+    lora_request = LoRARequest(
+        lora_name="befm",
+        lora_int_id=1,
+        lora_path=PEFT_MODEL_ID,
+    )
+    print(f"[INFO] PEFT adapter prepared: {PEFT_MODEL_ID}")
+
+    return llm, lora_request
+
+# Lazy load model and tokenizer
+_llm = None
+_lora_request = None
+_tokenizer = None
+
+def get_model_and_tokenizer():
+    global _llm, _lora_request, _tokenizer
+    if _llm is None:
+        _llm, _lora_request = load_model()
+        _tokenizer = AutoTokenizer.from_pretrained(BASE_MODEL_ID, token=HF_TOKEN)
+    return _llm, _lora_request, _tokenizer
+
+@spaces.GPU
+def generate_response(messages, max_new_tokens=512, temperature=0.7, top_p=0.9) -> str:
+    llm, lora_request, tokenizer = get_model_and_tokenizer()
+
+    # Apply Llama 3.1 chat template
+    prompt = tokenizer.apply_chat_template(
+        messages,
+        tokenize=False,
+        add_generation_prompt=True
+    )
+
+    sampling_params = SamplingParams(
+        temperature=temperature,
+        top_p=top_p,
+        max_tokens=max_new_tokens,
+    )
+
+    # Generate with vLLM
+    outputs = llm.generate(
+        prompts=[prompt],
+        sampling_params=sampling_params,
+        lora_request=lora_request,
+    )
+
+    return outputs[0].outputs[0].text
+
+def chat_fn(message, history, system_prompt, max_new_tokens, temperature, top_p):
+    # Build conversation in Llama 3.1 chat format
+    messages = []
+    if system_prompt:
+        messages.append({"role": "system", "content": system_prompt})
+
+    # History is already in dict format: [{"role": "user", "content": "..."}, ...]
+    for msg in (history or []):
+        messages.append(msg)
+
+    if message:
+        messages.append({"role": "user", "content": message})
+
+    reply = generate_response(
+        messages,
+        max_new_tokens=max_new_tokens,
+        temperature=temperature,
+        top_p=top_p,
+    )
+    return reply
+
+demo = gr.ChatInterface(
+    fn=lambda message, history, system_prompt, max_new_tokens, temperature, top_p:
+        chat_fn(message, history, system_prompt, max_new_tokens, temperature, top_p),
+    additional_inputs=[
+        gr.Textbox(label="System prompt (optional)", placeholder="You are Be.FM assistant...", lines=2),
+        gr.Slider(16, 2048, value=512, step=16, label="max_new_tokens"),
+        gr.Slider(0.1, 1.5, value=0.7, step=0.05, label="temperature"),
+        gr.Slider(0.1, 1.0, value=0.9, step=0.05, label="top_p"),
+    ],
+    title="Be.FM-8B (vLLM) - Fast Inference",
+    description="Chat interface using vLLM for optimized inference with Meta-Llama-3.1-8B-Instruct and PEFT adapter befm/Be.FM-8B."
+)
+
+if __name__ == "__main__":
+    demo.launch()

requirements.txt

@@ -3,3 +3,4 @@ transformers>=4.30.0
 peft>=0.4.0
 spaces
 accelerate
+vllm>=0.6.0