Upload folder using huggingface_hub

demo.py

@@ -1,7 +1,7 @@
 import gradio as gr
-from transformers import AutoModelForCausalLM, AutoTokenizer
-import torch
+from llama_cpp import Llama
 import logging
+import os
 
 # Setup logging
 logging.basicConfig(level=logging.INFO)
@@ -36,55 +36,42 @@ 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."""
     logger.info("Loading model...")
-    model = AutoModelForCausalLM.from_pretrained(
-        "Zwounds/boolean-search-model",
-        torch_dtype=torch.float32
+    model = Llama.from_pretrained(
+        repo_id="Zwounds/boolean-search-model",
+        filename="boolean.gguf",
     )
-    tokenizer = AutoTokenizer.from_pretrained("Zwounds/boolean-search-model")
-    tokenizer.use_default_system_prompt = False
     logger.info("Model loaded successfully")
-    
-    return model, tokenizer
+    return model
 
 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()
-    
+    if not output:
+        return ""
+    # Return the generated text, trimming any system prompts
     return output.strip()
 
-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)
+    prompt = f"""<|im_start|>system
+{SYSTEM_INSTRUCTION}<|im_end|>
+<|im_start|>user
+{query}<|im_end|>
+<|im_start|>assistant
+"""
     
     # Generate response
-    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
+    output = model.create_completion(
+        prompt,
+        max_tokens=64,
+        stop=["<|im_end|>"],
+        echo=False,
+        temperature=0.0
     )
     
-    return extract_response(tokenizer.batch_decode(outputs)[0])
+    return extract_response(output['choices'][0]['text'])
 
 # Example queries demonstrating various cases
 examples = [
@@ -136,7 +123,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"
@@ -152,7 +139,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,4 +1,2 @@
 gradio>=4.0.0
-huggingface-hub>=0.19.4
-transformers>=4.11.3
-torch>=1.9.0
+llama-cpp-python>=0.2.11