File size: 3,885 Bytes
e2ca773 506d5f3 e2ca773 83a0536 a34fd06 99c0ca6 506d5f3 e2ca773 a34fd06 e2ca773 99c0ca6 64d28ca 99c0ca6 83a0536 e2ca773 c6ebbc4 e2ca773 8a7855a e2ca773 8a7855a e2ca773 8a7855a e2ca773 506d5f3 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 |
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 # 1 KB
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)
# 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)
from llama_cpp import Llama
# Define the model URL and path
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"
# Download the model if it doesn't exist
if not os.path.exists(model_path):
print(f"Downloading model to {model_path}...")
download_model(model_url, model_path)
print("Model downloaded successfully.")
# Initialize the Llama model
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:"
# Generate response using Llama model
response = llm(full_prompt, max_tokens=150, temperature=0.7, top_p=0.9)
# Extract the generated recommendations
recommendations = response['choices'][0]['text'].strip()
return recommendations
# Create Gradio interface
iface = gr.Interface(
fn=generate_response,
inputs=gr.Textbox(lines=2, placeholder="Enter your query about LeetCode problems..."),
outputs="text",
title="LeetCode Problem Assistant",
description="Ask about LeetCode problems and get structured responses."
)
# Launch the app
iface.launch(share=True)
|