Upload folder using huggingface_hub

demo.py

@@ -1,6 +1,5 @@
 import gradio as gr
-from transformers import AutoModelForCausalLM, AutoTokenizer
-import torch
+from vllm import LLM
 import logging
 
 # Setup logging
@@ -36,56 +35,31 @@ SYSTEM_INSTRUCTION = """Convert natural language queries into boolean search que
    - Use OR with parentheses for alternatives"""
 
 def load_model():
-    """Load the model and set up tokenizer."""
+    """Load the model using vLLM."""
     logger.info("Loading model...")
-    model = AutoModelForCausalLM.from_pretrained(
-        "Zwounds/boolean-search-model",
-        device_map="cpu",
-        torch_dtype=torch.float32
+    model = LLM(
+        model="Zwounds/boolean-search-model",
+        tensor_parallel_size=1  # For CPU
     )
-    tokenizer = AutoTokenizer.from_pretrained("Zwounds/boolean-search-model")
-    tokenizer.use_default_system_prompt = False
     logger.info("Model loaded successfully")
-    
-    return model, tokenizer
-
-def extract_response(output: str) -> str:
-    """Extract the response part from the output."""
-    start_marker = "<|start_header_id|>assistant<|end_header_id|>"
-    end_marker = "<|eot_id|>"
-    
-    start_idx = output.find(start_marker)
-    if start_idx != -1:
-        start_idx += len(start_marker)
-        end_idx = output.find(end_marker, start_idx)
-        if end_idx != -1:
-            return output[start_idx:end_idx].strip()
-    
-    return output.strip()
+    return model
 
-def get_boolean_query(query: str, model=None, tokenizer=None) -> str:
+def get_boolean_query(query: str, model=None) -> str:
     """Generate boolean query from natural language."""
-    # Format the conversation
-    conversation = [
-        {"role": "system", "content": SYSTEM_INSTRUCTION},
-        {"role": "user", "content": query}
-    ]
-    
-    # Format into chat template
-    prompt = tokenizer.apply_chat_template(conversation, tokenize=False)
-    inputs = tokenizer(prompt, return_tensors="pt").to(model.device)
+    # Format the conversation with proper markers
+    prompt = f"""<|start_header_id|>system<|end_header_id|>{SYSTEM_INSTRUCTION}<|start_header_id|>user<|end_header_id|>{query}<|start_header_id|>assistant<|end_header_id|>"""
     
-    # Generate response
+    # Generate with vllm
     outputs = model.generate(
-        **inputs,
-        max_new_tokens=64,
-        do_sample=False,
-        use_cache=True,
-        pad_token_id=tokenizer.pad_token_id,
-        eos_token_id=tokenizer.eos_token_id
+        prompt,
+        max_tokens=64,
+        temperature=0.0,  # Deterministic
+        stop_tokens=["<|eot_id|>"]
     )
     
-    return extract_response(tokenizer.batch_decode(outputs)[0])
+    # Extract response
+    response = outputs[0].outputs[0].text.strip()
+    return response
 
 # Example queries demonstrating various cases
 examples = [
@@ -137,7 +111,7 @@ examples = [
 
 # Load model globally
 logger.info("Initializing model...")
-model, tokenizer = load_model()
+model = load_model()
 
 # Create Gradio interface
 title = "Natural Language to Boolean Search"
@@ -153,7 +127,7 @@ description = """Convert natural language queries into boolean search expression
 """
 
 demo = gr.Interface(
-    fn=lambda x: get_boolean_query(x, model, tokenizer),
+    fn=lambda x: get_boolean_query(x, model),
     inputs=[
         gr.Textbox(
             label="Enter your natural language query",

requirements.txt

@@ -1,5 +1,3 @@
 gradio>=4.0.0
-transformers>=4.0.0
-torch>=1.0.0
+vllm>=0.3.0
 huggingface-hub>=0.19.4
-accelerate>=0.26.0