File size: 3,039 Bytes
d186f2c |
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 |
from langchain.embeddings import HuggingFaceEmbeddings
from langchain.vectorstores import FAISS
from langchain.text_splitter import RecursiveCharacterTextSplitter
import streamlit as st
class RAGSystem:
def __init__(self):
self.embeddings = HuggingFaceEmbeddings(
model_name="sentence-transformers/all-mpnet-base-v2"
)
self.vector_store = None
self.text_splitter = RecursiveCharacterTextSplitter(
chunk_size=500,
chunk_overlap=50
)
def initialize_knowledge_base(self, knowledge_base):
"""Initialize vector store with construction knowledge"""
documents = []
for damage_type, cases in knowledge_base.items():
for case in cases:
doc_text = f"""
Damage Type: {damage_type}
Severity: {case['severity']}
Description: {case['description']}
Repair Methods: {', '.join(case['repair_method'])}
Cost: {case['estimated_cost']}
Timeframe: {case['timeframe']}
Location: {case['location']}
Required Expertise: {case['required_expertise']}
Immediate Action: {case['immediate_action']}
Prevention: {case['prevention']}
"""
documents.append(doc_text)
splits = self.text_splitter.create_documents(documents)
self.vector_store = FAISS.from_documents(splits, self.embeddings)
def query_knowledge_base(self, query, k=3):
"""Query the vector store for relevant information"""
if not self.vector_store:
raise ValueError("Vector store not initialized")
results = self.vector_store.similarity_search(query, k=k)
return results
def get_enhanced_analysis(self, damage_type, confidence, custom_query=None):
"""Get enhanced analysis based on damage type and confidence"""
if not custom_query:
query = f"What are the key considerations and recommendations for {damage_type} with {confidence}% confidence?"
else:
query = custom_query
relevant_docs = self.query_knowledge_base(query)
enhanced_info = {
"technical_details": [],
"safety_considerations": [],
"expert_recommendations": []
}
for doc in relevant_docs:
content = doc.page_content
if "severity" in content.lower() or "description" in content.lower():
enhanced_info["technical_details"].append(content)
elif "immediate action" in content.lower() or "prevention" in content.lower():
enhanced_info["safety_considerations"].append(content)
elif "repair" in content.lower() or "expertise" in content.lower():
enhanced_info["expert_recommendations"].append(content)
return enhanced_info |