|
import gradio as gr |
|
from pinecone import Pinecone |
|
import torch |
|
from pinecone_text.sparse import SpladeEncoder |
|
from sentence_transformers import SentenceTransformer |
|
from transformers import AutoTokenizer, AutoModelForCausalLM |
|
import os |
|
|
|
import requests |
|
import os |
|
from tqdm import tqdm |
|
|
|
|
|
def download_model(url, model_path): |
|
response = requests.get(url, stream=True) |
|
total_size = int(response.headers.get('content-length', 0)) |
|
block_size = 1024 |
|
|
|
with open(model_path, 'wb') as file, tqdm( |
|
desc=model_path, |
|
total=total_size, |
|
unit='iB', |
|
unit_scale=True, |
|
unit_divisor=1024, |
|
) as progress_bar: |
|
for data in response.iter_content(block_size): |
|
size = file.write(data) |
|
progress_bar.update(size) |
|
|
|
|
|
|
|
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) |
|
|
|
|
|
device = 'cpu' |
|
splade = SpladeEncoder(device=device) |
|
dense_model = SentenceTransformer('sentence-transformers/all-Mpnet-base-v2', device=device) |
|
|
|
from llama_cpp import Llama |
|
|
|
|
|
model_url = "https://huggingface.co/TheBloke/Llama-2-7B-Chat-GGUF/resolve/main/llama-2-7b-chat.Q4_K_M.gguf" |
|
model_path = "/tmp/llama-2-7b-chat.Q4_K_M.gguf" |
|
|
|
|
|
if not os.path.exists(model_path): |
|
print(f"Downloading model to {model_path}...") |
|
download_model(model_url, model_path) |
|
print("Model downloaded successfully.") |
|
|
|
|
|
llm = Llama(model_path=model_path, n_ctx=1024, n_threads=8, n_batch=512, |
|
mlock=True, n_gpu_layers=-1) |
|
|
|
|
|
def search_problems(query, top_k=5): |
|
dense_query = dense_model.encode([query])[0].tolist() |
|
sparse_query = splade.encode_documents([query])[0] |
|
|
|
results = index.query( |
|
vector=dense_query, |
|
sparse_vector={ |
|
'indices': sparse_query['indices'], |
|
'values': sparse_query['values'] |
|
}, |
|
top_k=top_k, |
|
include_metadata=True, |
|
namespace='leetcode-problems' |
|
) |
|
|
|
return results['matches'] |
|
|
|
def generate_few_shot_prompt(search_results): |
|
prompt = "Here are some example LeetCode problems:\n\n" |
|
for result in search_results: |
|
metadata = result['metadata'] |
|
prompt += f"Title: {metadata['title']}\n" |
|
prompt += f"Topics: {', '.join(metadata['topicTags'])}\n" |
|
prompt += f"Difficulty: {metadata['difficulty']}\n\n" |
|
return prompt |
|
|
|
def generate_response(user_query, top_k=5): |
|
search_results = search_problems(user_query, top_k) |
|
few_shot_prompt = generate_few_shot_prompt(search_results) |
|
|
|
system_prompt = """You are an AI assistant specialized in providing information about LeetCode problems. |
|
Your task is to recommend relevant problems based on the user's query and the provided examples. |
|
Focus on problem titles, difficulty levels, topic tags, and companies that have asked these problems. |
|
Do not provide specific problem solutions or content.""" |
|
|
|
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:" |
|
|
|
|
|
response = llm(full_prompt, max_tokens=150, temperature=0.7, top_p=0.9) |
|
|
|
|
|
recommendations = response['choices'][0]['text'].strip() |
|
return recommendations |
|
|
|
|
|
with gr.Blocks() as iface: |
|
gr.Markdown("# LeetCode Problem Assistant") |
|
|
|
disclaimer = gr.Textbox( |
|
value="• This is purely for Educational purpose only. We don't claim ownership of Problem statement\n" |
|
"• Copyright of problems and contents goes to Leetcode. Leetcode was not scraped as per terms.\n" |
|
"• Do not rely solely on this tool for interview preparation.\n" |
|
"• Always verify information with official LeetCode resources.\n" |
|
"• You can ask => 'I have upcoming interview with Google, I like to get better with Dynamic Programming. What do you recommend?'.\n" |
|
"• Your missing natural language based hybrid search on sparse and dense vectors built on Leetcode Knowledge Graph.", |
|
label="Disclaimer", |
|
interactive=False |
|
) |
|
|
|
query_input = gr.Textbox( |
|
lines=2, |
|
placeholder="Enter your natural language query about LeetCode problems...", |
|
label="Your Query" |
|
) |
|
|
|
output = gr.Textbox(label="Search Results") |
|
|
|
submit_button = gr.Button("Search") |
|
|
|
submit_button.click( |
|
fn=generate_response, |
|
inputs=query_input, |
|
outputs=output |
|
) |
|
|
|
|
|
iface.launch(share=True) |