File size: 11,332 Bytes
02acac5 |
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 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 199 200 201 202 203 204 205 206 207 208 209 210 211 212 213 214 215 216 217 218 219 220 221 222 223 224 225 226 227 228 229 230 231 232 233 234 235 236 237 238 239 240 241 242 243 244 245 246 247 248 249 250 251 252 253 254 255 256 257 258 259 260 261 262 263 264 265 266 267 268 269 270 271 272 273 274 275 276 277 278 279 280 281 282 283 284 285 286 287 288 289 290 291 292 293 294 295 296 297 298 299 300 301 302 303 304 305 306 307 308 309 310 311 312 313 314 315 316 317 318 319 320 321 322 323 324 325 326 327 328 329 330 331 332 333 334 335 336 337 338 339 340 341 342 343 344 345 346 347 348 349 350 351 352 353 354 355 356 |
import os
import json
import numpy as np
import faiss
import google.generativeai as genai
from typing import List, Dict, Any, Tuple
# Configure Gemini
API_KEY = os.getenv("GOOGLE_API_KEY")
if not API_KEY:
raise ValueError("Error: Set GOOGLE_API_KEY environment variable before running.")
genai.configure(api_key=API_KEY)
# File paths
DATA_DIR = "data"
PROFILE_EMBEDDINGS = os.path.join(DATA_DIR, "embeddings_profiles.jsonl")
JOB_EMBEDDINGS = os.path.join(DATA_DIR, "embeddings_jobs.jsonl")
# Global variables to cache loaded data
_profile_data = None
_job_data = None
_profile_index = None
_job_index = None
def load_embeddings_data(file_path: str) -> List[Dict[str, Any]]:
"""
Load embeddings data from JSONL file.
Args:
file_path: Path to the JSONL file containing embeddings
Returns:
List of dictionaries containing id, text, and embedding
"""
if not os.path.exists(file_path):
print(f"Warning: Embeddings file not found: {file_path}")
return []
data = []
try:
with open(file_path, 'r', encoding='utf-8') as f:
for line_num, line in enumerate(f, 1):
try:
record = json.loads(line.strip())
if 'embedding' in record and 'text' in record and 'id' in record:
data.append(record)
else:
print(f"Warning: Missing required fields in line {line_num} of {file_path}")
except json.JSONDecodeError as e:
print(f"Warning: JSON decode error in line {line_num} of {file_path}: {e}")
continue
except Exception as e:
print(f"Error reading {file_path}: {e}")
return []
return data
def build_faiss_index(embeddings: List[List[float]]) -> faiss.Index:
"""
Build a FAISS index from embeddings.
Args:
embeddings: List of embedding vectors
Returns:
FAISS index for similarity search
"""
if not embeddings:
return None
# Convert to numpy array
embedding_matrix = np.array(embeddings, dtype=np.float32)
# Get dimension
dimension = embedding_matrix.shape[1]
# Create FAISS index (using L2 distance)
index = faiss.IndexFlatL2(dimension)
# Add embeddings to index
index.add(embedding_matrix)
return index
def get_query_embedding(query: str, model: str = "models/text-embedding-004") -> List[float]:
"""
Get embedding for a query string using Gemini.
Args:
query: Query string to embed
model: Embedding model to use
Returns:
Embedding vector as list of floats
"""
try:
response = genai.embed_content(
model=model,
content=query,
task_type="retrieval_query" # Use query task type for search queries
)
return response['embedding']
except Exception as e:
print(f"Error getting query embedding: {e}")
# Return a zero vector as fallback (this won't give good results but prevents crashes)
return [0.0] * 768 # Assuming 768-dimensional embeddings
def initialize_profile_data():
"""Initialize profile data and FAISS index if not already loaded."""
global _profile_data, _profile_index
if _profile_data is None:
print("Loading profile embeddings...")
_profile_data = load_embeddings_data(PROFILE_EMBEDDINGS)
if _profile_data:
embeddings = [record['embedding'] for record in _profile_data]
_profile_index = build_faiss_index(embeddings)
print(f"Loaded {len(_profile_data)} profile embeddings")
else:
print("No profile embeddings found")
_profile_index = None
def initialize_job_data():
"""Initialize job data and FAISS index if not already loaded."""
global _job_data, _job_index
if _job_data is None:
print("Loading job embeddings...")
_job_data = load_embeddings_data(JOB_EMBEDDINGS)
if _job_data:
embeddings = [record['embedding'] for record in _job_data]
_job_index = build_faiss_index(embeddings)
print(f"Loaded {len(_job_data)} job embeddings")
else:
print("No job embeddings found")
_job_index = None
def retrieve_profiles(query: str, top_k: int = 5) -> List[Dict[str, Any]]:
"""
Retrieve the most similar profiles based on a query.
Args:
query: Search query string
top_k: Number of top results to return
Returns:
List of dictionaries containing profile information
"""
# Initialize data if needed
initialize_profile_data()
if not _profile_data or _profile_index is None:
print("No profile data available for search")
return []
# Get query embedding
query_embedding = get_query_embedding(query)
if not query_embedding:
return []
# Search using FAISS
query_vector = np.array([query_embedding], dtype=np.float32)
try:
# Search for top_k similar profiles
distances, indices = _profile_index.search(query_vector, min(top_k, len(_profile_data)))
# Prepare results
results = []
for i, idx in enumerate(indices[0]):
if idx < len(_profile_data): # Ensure valid index
profile = _profile_data[idx].copy()
profile['similarity_score'] = float(distances[0][i]) # Lower is better for L2 distance
results.append(profile)
return results
except Exception as e:
print(f"Error during profile search: {e}")
return []
def retrieve_jobs(query: str, top_k: int = 5) -> List[Dict[str, Any]]:
"""
Retrieve the most similar job listings based on a query.
Args:
query: Search query string
top_k: Number of top results to return
Returns:
List of dictionaries containing job information
"""
# Initialize data if needed
initialize_job_data()
if not _job_data or _job_index is None:
print("No job data available for search")
return []
# Get query embedding
query_embedding = get_query_embedding(query)
if not query_embedding:
return []
# Search using FAISS
query_vector = np.array([query_embedding], dtype=np.float32)
try:
# Search for top_k similar jobs
distances, indices = _job_index.search(query_vector, min(top_k, len(_job_data)))
# Prepare results
results = []
for i, idx in enumerate(indices[0]):
if idx < len(_job_data): # Ensure valid index
job = _job_data[idx].copy()
job['similarity_score'] = float(distances[0][i]) # Lower is better for L2 distance
results.append(job)
return results
except Exception as e:
print(f"Error during job search: {e}")
return []
def search_profiles_by_keywords(keywords: List[str], top_k: int = 5) -> List[Dict[str, Any]]:
"""
Search profiles using keyword matching as a fallback method.
Args:
keywords: List of keywords to search for
top_k: Number of top results to return
Returns:
List of matching profiles
"""
initialize_profile_data()
if not _profile_data:
return []
results = []
keywords_lower = [kw.lower() for kw in keywords]
for profile in _profile_data:
text_lower = profile['text'].lower()
score = sum(1 for kw in keywords_lower if kw in text_lower)
if score > 0:
profile_copy = profile.copy()
profile_copy['keyword_score'] = score
results.append(profile_copy)
# Sort by keyword score (descending)
results.sort(key=lambda x: x['keyword_score'], reverse=True)
return results[:top_k]
def search_jobs_by_keywords(keywords: List[str], top_k: int = 5) -> List[Dict[str, Any]]:
"""
Search jobs using keyword matching as a fallback method.
Args:
keywords: List of keywords to search for
top_k: Number of top results to return
Returns:
List of matching jobs
"""
initialize_job_data()
if not _job_data:
return []
results = []
keywords_lower = [kw.lower() for kw in keywords]
for job in _job_data:
text_lower = job['text'].lower()
score = sum(1 for kw in keywords_lower if kw in text_lower)
if score > 0:
job_copy = job.copy()
job_copy['keyword_score'] = score
results.append(job_copy)
# Sort by keyword score (descending)
results.sort(key=lambda x: x['keyword_score'], reverse=True)
return results[:top_k]
def get_stats():
"""Get statistics about loaded data."""
initialize_profile_data()
initialize_job_data()
profile_count = len(_profile_data) if _profile_data else 0
job_count = len(_job_data) if _job_data else 0
return {
"profiles_loaded": profile_count,
"jobs_loaded": job_count,
"profile_index_ready": _profile_index is not None,
"job_index_ready": _job_index is not None
}
# Test functions
def test_profile_search():
"""Test profile search functionality."""
test_queries = [
"Python developer",
"React frontend engineer",
"Data scientist with machine learning",
"Remote backend developer"
]
print("Testing profile search...")
for query in test_queries:
print(f"\nQuery: '{query}'")
results = retrieve_profiles(query, top_k=3)
for i, result in enumerate(results, 1):
print(f" {i}. {result['text'][:100]}...")
def test_job_search():
"""Test job search functionality."""
test_queries = [
"Remote React developer",
"Python backend engineer",
"Data science position",
"Full stack developer"
]
print("Testing job search...")
for query in test_queries:
print(f"\nQuery: '{query}'")
results = retrieve_jobs(query, top_k=3)
for i, result in enumerate(results, 1):
print(f" {i}. {result['text'][:100]}...")
if __name__ == "__main__":
# Print stats
stats = get_stats()
print("RAG Utils Statistics:")
print(f" Profiles loaded: {stats['profiles_loaded']}")
print(f" Jobs loaded: {stats['jobs_loaded']}")
print(f" Profile index ready: {stats['profile_index_ready']}")
print(f" Job index ready: {stats['job_index_ready']}")
# Run tests if data is available
if stats['profiles_loaded'] > 0:
test_profile_search()
if stats['jobs_loaded'] > 0:
test_job_search() |