Adding a simple monkey search for Leetcode - Darn LeetMonkey

app.py

@@ -1,11 +1,12 @@
 import gradio as gr
 from pinecone import Pinecone, ServerlessSpec
-from transformers import AutoTokenizer, AutoModelForCausalLM
 import torch
 from pinecone_text.sparse import SpladeEncoder
 from sentence_transformers import SentenceTransformer
 import transformers
 transformers.logging.set_verbosity_error()
+from transformers import AutoTokenizer, AutoModelForCausalLM, AutoGPTQConfig
+
 
 import os
 PINECONE_API_KEY = os.environ.get('PINECONE_API_KEY')
@@ -15,6 +16,11 @@ pc = Pinecone(api_key=PINECONE_API_KEY)
 index_name = "leetmonkey-sparse-dense"
 index = pc.Index(index_name)
 
+quantization_config = AutoGPTQConfig(
+    disable_exllama=True
+)
+
+
 # Initialize models
 device = 'cpu'
 splade = SpladeEncoder(device=device)
@@ -23,7 +29,7 @@ dense_model = SentenceTransformer('sentence-transformers/all-Mpnet-base-v2', dev
 # Load the quantized Llama 2 model and tokenizer
 model_name = "TheBloke/Llama-2-7B-Chat-GPTQ"
 tokenizer = AutoTokenizer.from_pretrained(model_name)
-model = AutoModelForCausalLM.from_pretrained(model_name, device_map="auto", quantization_config={"disable_exllama": True})
+model = AutoModelForCausalLM.from_pretrained(model_name, device_map="auto", quantization_config=quantization_config)
 
 def search_problems(query, top_k=5):
     dense_query = dense_model.encode([query])[0].tolist()