flutterbasit/RAG_App

1import os2import fitz  # For PDF extraction3from sentence_transformers import SentenceTransformer4import faiss5import numpy as np6from dotenv import load_dotenv7import streamlit as st8from groq import Groq9 10 11# Function to extract text from a PDF12def extract_text_from_pdf(file):13    try:14        doc = fitz.open(stream=file.read(), filetype="pdf")15        text = ""16        for page in doc:17            text += page.get_text()18        return text19    except Exception as e:20        st.error(f"Error extracting text: {e}")21        return ""22 23# Function to chunk the text24def chunk_text(text, chunk_size=500):25    sentences = text.split(". ")26    chunks = []27    current_chunk = ""28    for sentence in sentences:29        if len(current_chunk) + len(sentence) <= chunk_size:30            current_chunk += sentence + ". "31        else:32            chunks.append(current_chunk.strip())33            current_chunk = sentence + ". "34    if current_chunk:35        chunks.append(current_chunk.strip())36    return chunks37 38# Load the embedding model39embedding_model = SentenceTransformer('all-MiniLM-L6-v2')40 41# Function to generate embeddings42def generate_embeddings(chunks):43    return embedding_model.encode(chunks)44 45# Function to store embeddings in FAISS46def store_embeddings_in_faiss(embeddings):47    try:48        dimension = embeddings.shape[1]49        index = faiss.IndexFlatL2(dimension)50        index.add(embeddings)51        return index52    except Exception as e:53        st.error(f"Error with FAISS: {e}")54        return None55 56# Function to retrieve similar chunks57def retrieve_similar_chunks(query, index, chunks, model):58    try:59        query_embedding = model.encode([query])[0]60        distances, indices = index.search(np.array([query_embedding]), k=5)61        return [chunks[i] for i in indices[0]]62    except Exception as e:63        st.error(f"Error retrieving similar chunks: {e}")64        return []65 66# Load environment variables67# Initialize Groq client with direct API key68groq_api_key = "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"69 70if not groq_api_key:71    st.error("The GROQ_API_KEY is not set.")72    exit()73 74# Initialize Groq client75groq_client = Groq(api_key=groq_api_key)76 77 78def query_llm(prompt, model="llama3-8b-8192"):79    try:80        response = groq_client.chat.completions.create(81            messages=[82                {"role": "system", "content": "You are a helpful assistant."},83                {"role": "user", "content": prompt},84            ],85            model=model,86        )87        return response.choices[0].message.content88    except Exception as e:89        st.error(f"Error querying LLM: {e}")90        return "Error in LLM response."91 92# Streamlit application93def main():94    st.title("RAG Application with Groq API")95    96    # File upload97    uploaded_file = st.file_uploader("Upload a PDF", type="pdf")98    if uploaded_file:99        # Extract text100        pdf_text = extract_text_from_pdf(uploaded_file)101        if not pdf_text:102            return103        104        st.write("PDF Text Extracted:")105        st.write(pdf_text[:500])  # Show a preview106        107        # Chunk the text108        chunks = chunk_text(pdf_text)109        st.write(f"Text split into {len(chunks)} chunks.")110        111        # Generate embeddings112        embeddings = np.array(generate_embeddings(chunks))113        index = store_embeddings_in_faiss(embeddings)114        if index is None:115            return116        117        # Query handling118        query = st.text_input("Enter your query:")119        if query:120            similar_chunks = retrieve_similar_chunks(query, index, chunks, embedding_model)121            st.write("Relevant Chunks:")122            for i, chunk in enumerate(similar_chunks, start=1):123                st.write(f"Chunk {i}: {chunk}")124 125            # Query the LLM126            combined_context = " ".join(similar_chunks[:3])127            llm_prompt = f"Context: {combined_context}\n\nQuery: {query}"128            llm_response = query_llm(llm_prompt)129            st.write("LLM Response:")130            st.write(llm_response)131 132if __name__ == "__main__":133    main()134