File size: 4,086 Bytes
c85d1d3 |
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 |
# rag.py
from sentence_transformers import SentenceTransformer
import faiss
import numpy as np
import google as genai
import os
class SimpleRAG:
def __init__(self, api_key):
# Initialize the embedding model and generative AI
self.embedder = SentenceTransformer('sentence-transformers/all-MiniLM-L6-v2')
genai.configure(api_key=api_key)
self.model = genai.GenerativeModel("gemini-1.5-flash")
self.index = None
self.chunks = []
self.is_initialized = False
self.processing_status = None
def chunk_text(self, text, chunk_size=700):
"""Split text into smaller chunks."""
words = text.split()
return [' '.join(words[i:i + chunk_size]) for i in range(0, len(words), chunk_size)]
def process_search_data(self, search_data):
"""
Process search result data and index it.
'search_data' should be a list of job posting dictionaries.
For each job posting, we combine key fields (e.g., job title and description) and then chunk the text.
"""
try:
self.processing_status = "Processing search data..."
combined_text = ""
for job in search_data:
# Combine job title and job description (you can add more fields if needed)
job_title = job.get('job_title', '')
job_description = job.get('job_description', '')
combined_text += f"Job Title: {job_title}. Description: {job_description}. "
if not combined_text.strip():
raise Exception("No text found in search results.")
# Chunk the combined text
self.chunks = self.chunk_text(combined_text)
if not self.chunks:
raise Exception("No content chunks were generated from search data.")
# Generate embeddings and create the FAISS index
embeddings = self.embedder.encode(self.chunks)
vector_dimension = embeddings.shape[1]
self.index = faiss.IndexFlatL2(vector_dimension)
self.index.add(np.array(embeddings).astype('float32'))
self.is_initialized = True
self.processing_status = f"RAG system initialized with {len(self.chunks)} chunks."
return {"status": "success", "message": self.processing_status}
except Exception as e:
self.processing_status = f"Error: {str(e)}"
self.is_initialized = False
return {"status": "error", "message": str(e)}
def get_status(self):
"""Return the current processing status."""
return {
"is_initialized": self.is_initialized,
"status": self.processing_status
}
def get_relevant_chunks(self, query, k=3):
"""Retrieve the top-k most relevant text chunks for a given query."""
query_vector = self.embedder.encode([query])
distances, chunk_indices = self.index.search(query_vector.astype('float32'), k)
return [self.chunks[i] for i in chunk_indices[0]]
def query(self, question):
"""Query the RAG system with a user question."""
if not self.is_initialized:
raise Exception("RAG system not initialized. Please process search data first.")
try:
context = self.get_relevant_chunks(question)
prompt = f"""
Based on the following context, provide a clear and concise answer.
If the context doesn't contain enough relevant information, say "I don't have enough information to answer that question."
Context:
{' '.join(context)}
Question: {question}
"""
response = self.model.generate_content(prompt)
return {
"status": "success",
"answer": response.text.strip(),
"context": context
}
except Exception as e:
return {
"status": "error",
"message": str(e)
}
|