File size: 14,344 Bytes
62d1e75
 
 
 
 
 
 
 
 
 
 
 
 
 
2311e2d
 
 
 
 
 
 
 
 
 
 
 
 
 
62d1e75
 
2311e2d
 
 
 
 
 
 
 
 
 
 
 
62d1e75
 
2311e2d
 
 
 
 
 
 
 
 
62d1e75
 
 
 
 
 
 
 
 
 
 
 
 
 
 
3b2ec72
 
ac3798b
 
 
 
 
 
 
 
 
 
 
 
3b2ec72
 
ac3798b
 
 
3b2ec72
ac3798b
 
 
 
3b2ec72
62d1e75
 
 
3b2ec72
 
62d1e75
2311e2d
62d1e75
38deecc
 
 
 
2311e2d
38deecc
 
 
2311e2d
 
 
 
 
 
 
 
 
 
 
87dead9
 
 
 
 
 
 
 
 
2311e2d
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
62d1e75
 
 
 
 
 
 
 
 
 
 
ff4d9c5
62d1e75
3b2ec72
 
 
 
ff4d9c5
3b2ec72
 
ff4d9c5
3b2ec72
62d1e75
ff4d9c5
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
62d1e75
2311e2d
62d1e75
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
38deecc
62d1e75
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
38deecc
62d1e75
 
38deecc
f637309
 
 
 
 
 
 
 
 
 
38deecc
62d1e75
 
 
2311e2d
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
62d1e75
 
 
 
 
 
 
 
38deecc
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
357
358
359
360
361
362
363
364
365
366
367
368
369
370
371
372
import os
import json
import numpy as np
import faiss
import torch
import torch.nn as nn
from google.cloud import storage
from transformers import AutoTokenizer, AutoModel
import openai
import textwrap
import unicodedata
import streamlit as st
from utils import setup_gcp_auth, setup_openai_auth

# Force model to CPU for stability
os.environ["CUDA_VISIBLE_DEVICES"] = ""

# Create a function to initialize session state
def initialize_session_state():
    if 'model_initialized' not in st.session_state:
        st.session_state.model_initialized = False
        st.session_state.model = None
        st.session_state.tokenizer = None
        st.session_state.device = torch.device("cpu")
        print("Initialized session state variables")

# Call the initialization function right away
initialize_session_state()

# Load GCP authentication from utility function
def setup_gcp_client():
    try:
        credentials = setup_gcp_auth()
        storage_client = storage.Client(credentials=credentials)
        bucket_name = "indian_spiritual-1"
        bucket = storage_client.bucket(bucket_name)
        print("βœ… GCP client initialized successfully")
        return bucket
    except Exception as e:
        print(f"❌ GCP client initialization error: {str(e)}")
        st.error(f"GCP client initialization error: {str(e)}")
        return None

# Setup OpenAI authentication
def setup_openai_client():
    try:
        setup_openai_auth()
        print("βœ… OpenAI client initialized successfully")
        return True
    except Exception as e:
        print(f"❌ OpenAI client initialization error: {str(e)}")
        st.error(f"OpenAI client initialization error: {str(e)}")
        return False

# GCS Paths
metadata_file_gcs = "metadata/metadata.jsonl"
embeddings_file_gcs = "processed/embeddings/all_embeddings.npy"
faiss_index_file_gcs = "processed/indices/faiss_index.faiss"
text_chunks_file_gcs = "processed/chunks/text_chunks.txt"

# Local Paths
local_embeddings_file = "all_embeddings.npy"
local_faiss_index_file = "faiss_index.faiss"
local_text_chunks_file = "text_chunks.txt"
local_metadata_file = "metadata.jsonl"

def load_model():
    try:
        # Initialize model if it doesn't exist
        if 'model' not in st.session_state or st.session_state.model is None:
            # Force model to CPU - more stable than GPU for this use case
            os.environ["CUDA_VISIBLE_DEVICES"] = ""
            
            print("Loading tokenizer...")
            tokenizer = AutoTokenizer.from_pretrained("intfloat/e5-small-v2")
            
            print("Loading model...")
            model = AutoModel.from_pretrained(
                "intfloat/e5-small-v2",
                torch_dtype=torch.float16  # Use half precision
            )
            
            # Move model to CPU explicitly
            model = model.to('cpu')
            model.eval()
            torch.set_grad_enabled(False)
            
            # Store in session state
            st.session_state.tokenizer = tokenizer
            st.session_state.model = model
            
            print("βœ… Model loaded successfully")
        
        return st.session_state.tokenizer, st.session_state.model
    except Exception as e:
        print(f"❌ Error loading model: {str(e)}")
        # Return None values instead of raising to avoid crashing
        return None, None

def download_file_from_gcs(bucket, gcs_path, local_path):
    """Download a file from GCS to local storage."""
    try:
        blob = bucket.blob(gcs_path)
        blob.download_to_filename(local_path)
        print(f"βœ… Downloaded {gcs_path} β†’ {local_path}")
        return True
    except Exception as e:
        print(f"❌ Error downloading {gcs_path}: {str(e)}")
        st.error(f"Error downloading {gcs_path}: {str(e)}")
        return False

def load_data_files():
    # Initialize GCP and OpenAI clients
    bucket = setup_gcp_client()
    openai_initialized = setup_openai_client()
    
    if not bucket or not openai_initialized:
        st.error("Failed to initialize required services")
        return None, None, None
    
    # Download necessary files - remove the spinner from here
    success = True
    success &= download_file_from_gcs(bucket, faiss_index_file_gcs, local_faiss_index_file)
    success &= download_file_from_gcs(bucket, text_chunks_file_gcs, local_text_chunks_file)
    success &= download_file_from_gcs(bucket, metadata_file_gcs, local_metadata_file)
    
    if not success:
        st.error("Failed to download required files")
        return None, None, None
    
    # Load FAISS index
    try:
        faiss_index = faiss.read_index(local_faiss_index_file)
    except Exception as e:
        print(f"❌ Error loading FAISS index: {str(e)}")
        st.error(f"Error loading FAISS index: {str(e)}")
        return None, None, None
    
    # Load text chunks
    try:
        text_chunks = {}  # {ID -> (Title, Author, Text)}
        with open(local_text_chunks_file, "r", encoding="utf-8") as f:
            for line in f:
                parts = line.strip().split("\t")
                if len(parts) == 4:
                    text_chunks[int(parts[0])] = (parts[1], parts[2], parts[3])
    except Exception as e:
        print(f"❌ Error loading text chunks: {str(e)}")
        st.error(f"Error loading text chunks: {str(e)}")
        return None, None, None
    
    # Load metadata.jsonl for publisher information
    try:
        metadata_dict = {}
        with open(local_metadata_file, "r", encoding="utf-8") as f:
            for line in f:
                item = json.loads(line)
                metadata_dict[item["Title"]] = item  # Store for easy lookup
    except Exception as e:
        print(f"❌ Error loading metadata: {str(e)}")
        st.error(f"Error loading metadata: {str(e)}")
        return None, None, None
    
    print(f"βœ… FAISS index and text chunks loaded. {len(text_chunks)} passages available.")
    return faiss_index, text_chunks, metadata_dict

def average_pool(last_hidden_states, attention_mask):
    """Average pooling for sentence embeddings."""
    last_hidden = last_hidden_states.masked_fill(~attention_mask[..., None].bool(), 0.0)
    return last_hidden.sum(dim=1) / attention_mask.sum(dim=1)[..., None]

query_embedding_cache = {}

def get_embedding(text):
    if text in query_embedding_cache:
        return query_embedding_cache[text]

    try:
        # Ensure model initialization
        if 'model' not in st.session_state or st.session_state.model is None:
            tokenizer, model = load_model()
            if model is None:
                return np.zeros((1, 384), dtype=np.float32)  # Fallback for e5-small-v2
        else:
            tokenizer, model = st.session_state.tokenizer, st.session_state.model
            
        input_text = f"query: {text}" if len(text) < 512 else f"passage: {text}"
        
        # Explicitly specify truncation parameters to avoid warnings
        inputs = tokenizer(
            input_text,
            padding=True,
            truncation=True,
            return_tensors="pt",
            max_length=512,
            return_attention_mask=True
        )
        
        # Move to CPU explicitly before processing
        inputs = {k: v.to('cpu') for k, v in inputs.items()}
        
        with torch.no_grad():
            outputs = model(**inputs)
            embeddings = average_pool(outputs.last_hidden_state, inputs['attention_mask'])
            embeddings = nn.functional.normalize(embeddings, p=2, dim=1)
            # Ensure we detach and move to numpy on CPU
            embeddings = embeddings.detach().cpu().numpy()
            
        # Explicitly clean up
        del outputs
        torch.cuda.empty_cache() if torch.cuda.is_available() else None
        
        query_embedding_cache[text] = embeddings
        return embeddings
    except Exception as e:
        print(f"❌ Embedding error: {str(e)}")
        st.error(f"Embedding error: {str(e)}")
        return np.zeros((1, 384), dtype=np.float32)  # Changed from 1024 to 384 for e5-small-v2

def retrieve_passages(query, faiss_index, text_chunks, metadata_dict, top_k=5, similarity_threshold=0.5):
    """Retrieve top-k most relevant passages using FAISS with metadata."""
    try:
        print(f"\nπŸ” Retrieving passages for query: {query}")
        query_embedding = get_embedding(query)
        distances, indices = faiss_index.search(query_embedding, top_k * 2)

        print(f"Found {len(distances[0])} potential matches")
        retrieved_passages = []
        retrieved_sources = []
        cited_titles = set()

        for dist, idx in zip(distances[0], indices[0]):
            print(f"Distance: {dist:.4f}, Index: {idx}")
            if idx in text_chunks and dist >= similarity_threshold:
                title_with_txt, author, text = text_chunks[idx]

                # Normalize title and remove .txt
                clean_title = title_with_txt.replace(".txt", "") if title_with_txt.endswith(".txt") else title_with_txt
                clean_title = unicodedata.normalize("NFC", clean_title)

                # Ensure unique citations
                if clean_title in cited_titles:
                    continue  

                metadata_entry = metadata_dict.get(clean_title, {})
                author = metadata_entry.get("Author", "Unknown")
                publisher = metadata_entry.get("Publisher", "Unknown")

                cited_titles.add(clean_title)

                retrieved_passages.append(text)
                retrieved_sources.append((clean_title, author, publisher))  

                if len(retrieved_passages) == top_k:
                    break

        print(f"Retrieved {len(retrieved_passages)} passages")
        return retrieved_passages, retrieved_sources
    except Exception as e:
        print(f"❌ Error in retrieve_passages: {str(e)}")
        st.error(f"Error in retrieve_passages: {str(e)}")
        return [], []

def answer_with_llm(query, context=None, word_limit=100):
    """
    Generate an answer using OpenAI GPT model with formatted citations.
    """
    try:
        if context:
            formatted_contexts = []
            total_chars = 0
            max_context_chars = 4000  

            for (title, author, publisher), text in context:
                remaining_space = max(0, max_context_chars - total_chars)
                excerpt_len = min(150, remaining_space)

                if excerpt_len > 50:
                    excerpt = text[:excerpt_len].strip() + "..." if len(text) > excerpt_len else text
                    formatted_context = f"[{title} by {author}, Published by {publisher}] {excerpt}"
                    formatted_contexts.append(formatted_context)
                    total_chars += len(formatted_context)

                if total_chars >= max_context_chars:
                    break

            formatted_context = "\n".join(formatted_contexts)
        else:
            formatted_context = "No relevant information available."

        # System message
        system_message = (
            "You are an AI specialized in Indian spiritual texts. "
            "Answer based on context, summarizing ideas rather than quoting verbatim. "
            "Ensure proper citation and do not include direct excerpts."
        )

        user_message = f"""
        Context:
        {formatted_context}
        Question:
        {query}
        """

        response = openai.chat.completions.create(
            model="gpt-3.5-turbo",
            messages=[
                {"role": "system", "content": system_message},
                {"role": "user", "content": user_message}
            ],
            max_tokens=200,
            temperature=0.7
        )

        answer = response.choices[0].message.content.strip()

        # Enforce word limit
        words = answer.split()
        if len(words) > word_limit:
            answer = " ".join(words[:word_limit])
            if not answer.endswith((".", "!", "?")):
                answer += "."

        return answer

    except Exception as e:
        print(f"❌ LLM API error: {str(e)}")
        st.error(f"LLM API error: {str(e)}")
        return "I apologize, but I'm unable to answer at the moment."

def format_citations(sources):
    """Format citations to display each one on a new line with a full stop if needed."""
    formatted_citations = []
    for title, author, publisher in sources:
        # Check if the publisher already ends with a period, question mark, or exclamation mark
        if publisher.endswith(('.', '!', '?')):
            formatted_citations.append(f"πŸ“š {title} by {author}, Published by {publisher}")
        else:
            formatted_citations.append(f"πŸ“š {title} by {author}, Published by {publisher}.")
    
    return "\n".join(formatted_citations)

def process_query(query, top_k=5, word_limit=100):
    """Process a query through the RAG pipeline with proper formatting."""
    print(f"\nπŸ” Processing query: {query}")
    
    # Load data files if not already loaded
    if not hasattr(st.session_state, 'data_loaded') or not st.session_state.data_loaded:
        st.session_state.faiss_index, st.session_state.text_chunks, st.session_state.metadata_dict = load_data_files()
        st.session_state.data_loaded = True
    
    # Check if data loaded successfully
    if not st.session_state.faiss_index or not st.session_state.text_chunks or not st.session_state.metadata_dict:
        return {"query": query, "answer_with_rag": "⚠️ System error: Data files not loaded properly.", "citations": "No citations available."}

    retrieved_context, retrieved_sources = retrieve_passages(
        query, 
        st.session_state.faiss_index, 
        st.session_state.text_chunks, 
        st.session_state.metadata_dict,
        top_k=top_k
    )
    
    sources = format_citations(retrieved_sources) if retrieved_sources else "No citation available."

    if retrieved_context:
        context_with_sources = list(zip(retrieved_sources, retrieved_context))
        llm_answer_with_rag = answer_with_llm(query, context_with_sources, word_limit=word_limit)
    else:
        llm_answer_with_rag = "⚠️ No relevant context found."

    return {"query": query, "answer_with_rag": llm_answer_with_rag, "citations": sources}