Upload folder using huggingface_hub

demo.py

@@ -1,5 +1,6 @@
 import gradio as gr
-from vllm import LLM
+from transformers import AutoModelForCausalLM, AutoTokenizer
+import torch
 import logging
 
 # Setup logging
@@ -35,31 +36,55 @@ SYSTEM_INSTRUCTION = """Convert natural language queries into boolean search que
    - Use OR with parentheses for alternatives"""
 
 def load_model():
-    """Load the model using vLLM."""
+    """Load the model and set up tokenizer."""
     logger.info("Loading model...")
-    model = LLM(
-        model="Zwounds/boolean-search-model",
-        tensor_parallel_size=1  # For CPU
+    model = AutoModelForCausalLM.from_pretrained(
+        "Zwounds/boolean-search-model",
+        torch_dtype=torch.float32
     )
+    tokenizer = AutoTokenizer.from_pretrained("Zwounds/boolean-search-model")
+    tokenizer.use_default_system_prompt = False
     logger.info("Model loaded successfully")
-    return model
+    
+    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()
 
-def get_boolean_query(query: str, model=None) -> str:
+def get_boolean_query(query: str, model=None, tokenizer=None) -> str:
     """Generate boolean query from natural language."""
-    # 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|>"""
+    # 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)
     
-    # Generate with vllm
+    # Generate response
     outputs = model.generate(
-        prompt,
-        max_tokens=64,
-        temperature=0.0,  # Deterministic
-        stop_tokens=["<|eot_id|>"]
+        **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
     )
     
-    # Extract response
-    response = outputs[0].outputs[0].text.strip()
-    return response
+    return extract_response(tokenizer.batch_decode(outputs)[0])
 
 # Example queries demonstrating various cases
 examples = [
@@ -111,7 +136,7 @@ examples = [
 
 # Load model globally
 logger.info("Initializing model...")
-model = load_model()
+model, tokenizer = load_model()
 
 # Create Gradio interface
 title = "Natural Language to Boolean Search"
@@ -127,7 +152,7 @@ description = """Convert natural language queries into boolean search expression
 """
 
 demo = gr.Interface(
-    fn=lambda x: get_boolean_query(x, model),
+    fn=lambda x: get_boolean_query(x, model, tokenizer),
     inputs=[
         gr.Textbox(
             label="Enter your natural language query",

requirements.txt

@@ -1,3 +1,3 @@
 gradio>=4.0.0
-vllm>=0.3.0
 huggingface-hub>=0.19.4
+transformers>=4.11.3