Adding a simple monkey search for Leetcode - Darn LeetMonkey

app.py

@@ -1,38 +1,27 @@
 import gradio as gr
-from pinecone import Pinecone, ServerlessSpec
+from pinecone import Pinecone
 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
-
-
 import os
-PINECONE_API_KEY = os.environ.get('PINECONE_API_KEY')
 
+# Initialize Pinecone
+PINECONE_API_KEY = os.environ.get('PINECONE_API_KEY')
 pc = Pinecone(api_key=PINECONE_API_KEY)
-
 index_name = "leetmonkey-sparse-dense"
 index = pc.Index(index_name)
 
-
-
 # Initialize models
 device = 'cpu'
 splade = SpladeEncoder(device=device)
 dense_model = SentenceTransformer('sentence-transformers/all-Mpnet-base-v2', device=device)
 
 # Load the quantized Llama 2 model and tokenizer
-model_name = "TheBloke/Llama-2-7B-Chat-GGML"
+model_name = "distilgpt2"  # Using distilgpt2 for CPU efficiency
 tokenizer = AutoTokenizer.from_pretrained(model_name)
 model = AutoModelForCausalLM.from_pretrained(model_name, low_cpu_mem_usage=True)
 
-
-# Disable Exllama backend if needed
-if hasattr(model, 'quantization_config'):
-    model.quantization_config.use_exllama = False
-
 def search_problems(query, top_k=5):
     dense_query = dense_model.encode([query])[0].tolist()
     sparse_query = splade.encode_documents([query])[0]
@@ -71,13 +60,14 @@ def generate_response(user_query, top_k=5):
     user_prompt = f"Based on the following query, recommend relevant LeetCode problems:\n{user_query}"
     full_prompt = f"{system_prompt}\n\n{few_shot_prompt}\n{user_prompt}\n\nRecommendations:"
 
-    input_ids = tokenizer.encode(full_prompt, return_tensors="pt").to(model.device)
+    input_ids = tokenizer.encode(full_prompt, return_tensors="pt").to(device)
     attention_mask = torch.ones_like(input_ids)
 
     with torch.no_grad():
         output = model.generate(
             input_ids,
-            max_new_tokens=100,  # Reduce this for faster generation
+            attention_mask=attention_mask,
+            max_new_tokens=100,  # Adjust as needed
             do_sample=True,
             top_p=0.9,
             temperature=0.7,
@@ -99,4 +89,4 @@ iface = gr.Interface(
 )
 
 # Launch the app
-iface.launch(share=True)
+iface.launch(share=True)