File size: 4,381 Bytes
8f69dec |
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 |
import os
import pandas as pd
import json
from langchain_community.vectorstores import FAISS
from langchain_huggingface import HuggingFaceEmbeddings
from langchain_text_splitters import RecursiveCharacterTextSplitter
from langchain_core.documents import Document
from transformers import pipeline
class RAGSystem:
def __init__(self, data_dir):
self.data_dir = data_dir
self.vector_store = None
self.embeddings = HuggingFaceEmbeddings(model_name="all-MiniLM-L6-v2")
self.llm_pipeline = None
def load_documents(self):
docs = []
# 1. Load Reviews (Unstructured)
reviews_path = os.path.join(self.data_dir, 'customer_reviews.csv')
if os.path.exists(reviews_path):
df_reviews = pd.read_csv(reviews_path)
for _, row in df_reviews.iterrows():
content = f"Product: {row.get('Product', 'Unknown')}\nDate: {row.get('Date', '')}\nRating: {row.get('Rating', '')}\nReview: {row.get('ReviewText', '')}"
metadata = {"source": "customer_reviews", "product": row.get('Product', 'Unknown')}
docs.append(Document(page_content=content, metadata=metadata))
# 2. Load Web Logs (Semi-structured)
logs_path = os.path.join(self.data_dir, 'web_logs.json')
if os.path.exists(logs_path):
with open(logs_path, 'r') as f:
logs_data = json.load(f)
for log in logs_data:
content = f"Log Timestamp: {log.get('timestamp', '')}\nAction: {log.get('action', '')}\nPage: {log.get('page', '')}\nUser: {log.get('user_id', '')}"
metadata = {"source": "web_logs"}
docs.append(Document(page_content=content, metadata=metadata))
# 3. Load Sales Summary (Structured -> Text)
sales_path = os.path.join(self.data_dir, 'sales_data.csv')
if os.path.exists(sales_path):
df_sales = pd.read_csv(sales_path)
# Create a summary per product/region instead of every row to save tokens/index size
summary = df_sales.groupby(['Product', 'Region'])['TotalPrice'].sum().reset_index()
for _, row in summary.iterrows():
content = f"Sales Summary:\nProduct: {row['Product']}\nRegion: {row['Region']}\nTotal Revenue: ${row['TotalPrice']:.2f}"
metadata = {"source": "sales_summary"}
docs.append(Document(page_content=content, metadata=metadata))
return docs
def build_index(self):
docs = self.load_documents()
text_splitter = RecursiveCharacterTextSplitter(chunk_size=500, chunk_overlap=50)
splits = text_splitter.split_documents(docs)
if splits:
self.vector_store = FAISS.from_documents(splits, self.embeddings)
return True
return False
def init_llm(self):
# Initialize a small local model for generation
# Using flan-t5-small as it is lightweight
try:
self.llm_pipeline = pipeline("text2text-generation", model="google/flan-t5-small")
except Exception as e:
print(f"Error loading LLM: {e}")
self.llm_pipeline = None
def query(self, user_query, k=3):
if not self.vector_store:
return {
"answer": "System not initialized. Please build the index first.",
"context": []
}
# Retrieve
docs = self.vector_store.similarity_search(user_query, k=k)
context_text = "\n\n".join([d.page_content for d in docs])
# Generate
if self.llm_pipeline:
prompt = f"Summarize the following context to answer the question. \n\nContext:\n{context_text}\n\nQuestion: {user_query}\n\nAnswer:"
# Truncate prompt if too long (simple heuristic)
if len(prompt) > 2048:
prompt = prompt[:2048]
result = self.llm_pipeline(prompt, max_length=200, do_sample=False)
answer = result[0]['generated_text']
else:
answer = "LLM not loaded. Displaying retrieved context only."
return {
"answer": answer,
"context": docs
}
|