File size: 3,611 Bytes
e2ca773 8fdb4d5 ea3e553 e2ca773 a34fd06 e2ca773 8fdb4d5 ac606ee 8fdb4d5 ea3e553 e2ca773 a34fd06 e2ca773 8fdb4d5 e2ca773 a34fd06 |
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 |
import gradio as gr
from pinecone import Pinecone, ServerlessSpec
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, GPTQConfig
import os
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)
quantization_config = GPTQConfig(
bits=8,
disable_exllama=True
)
# 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-GPTQ"
tokenizer = AutoTokenizer.from_pretrained(model_name)
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()
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:"
input_ids = tokenizer.encode(full_prompt, return_tensors="pt").to(model.device)
attention_mask = torch.ones_like(input_ids)
with torch.no_grad():
output = model.generate(
input_ids,
attention_mask=attention_mask,
max_new_tokens=250,
do_sample=True,
top_p=0.9,
temperature=0.7,
num_return_sequences=1,
pad_token_id=tokenizer.eos_token_id
)
response = tokenizer.decode(output[0], skip_special_tokens=True)
recommendations = response.split("Recommendations:")[1].strip()
return recommendations
# Create a Gradio interface
iface = gr.Interface(
fn=generate_response,
inputs=gr.inputs.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 based on titles, topics, and difficulty levels."
)
# Launch the app
iface.launch(share=True) |