Fix Gradio ChatInterface: remove lambda wrapper, add lazy loading, make public

app.py

@@ -65,12 +65,22 @@ def load_model_and_tokenizer():
             return base, tok
     return base, tok
 
-model, tokenizer = load_model_and_tokenizer()
-DEVICE = model.device
+# Lazy load model and tokenizer
+_model = None
+_tokenizer = None
+
+def get_model_and_tokenizer():
+    global _model, _tokenizer
+    if _model is None:
+        _model, _tokenizer = load_model_and_tokenizer()
+    return _model, _tokenizer
 
 @spaces.GPU
 @torch.inference_mode()
 def generate_response(messages, max_new_tokens=512, temperature=0.7, top_p=0.9) -> str:
+    model, tokenizer = get_model_and_tokenizer()
+    device = model.device
+
     # Apply Llama 3.1 chat template
     prompt = tokenizer.apply_chat_template(
         messages,
@@ -78,7 +88,7 @@ def generate_response(messages, max_new_tokens=512, temperature=0.7, top_p=0.9)
         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()}
+    enc = {k: v.to(device) for k, v in enc.items()}
 
     input_length = enc['input_ids'].shape[1]
     out = model.generate(
@@ -90,7 +100,8 @@ def generate_response(messages, max_new_tokens=512, temperature=0.7, top_p=0.9)
         pad_token_id=tokenizer.eos_token_id,
     )
     # Decode only the newly generated tokens
-    return tokenizer.decode(out[0][input_length:], skip_special_tokens=True)
+    generated_text = tokenizer.decode(out[0][input_length:], skip_special_tokens=True)
+    return generated_text.strip()
 
 def chat_fn(message, history, system_prompt, max_new_tokens, temperature, top_p):
     # Build conversation in Llama 3.1 chat format
@@ -117,8 +128,7 @@ def chat_fn(message, history, system_prompt, max_new_tokens, temperature, 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),
+    fn=chat_fn,
     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"),
@@ -130,4 +140,4 @@ demo = gr.ChatInterface(
 )
 
 if __name__ == "__main__":
-    demo.launch()
+    demo.launch(share=True)