DrishtiSharma commited on
Commit
08e5756
Β·
verified Β·
1 Parent(s): 6239518

Delete lab/test2.py

Browse files
Files changed (1) hide show
  1. lab/test2.py +0 -135
lab/test2.py DELETED
@@ -1,135 +0,0 @@
1
- import os
2
- import chromadb
3
- import requests
4
- import streamlit as st
5
- from langchain.chains import LLMChain
6
- from langchain.prompts import PromptTemplate
7
- from langchain_groq import ChatGroq
8
- from langchain.document_loaders import PDFPlumberLoader
9
- from langchain_experimental.text_splitter import SemanticChunker
10
- from langchain_huggingface import HuggingFaceEmbeddings
11
- from langchain_chroma import Chroma
12
- from prompts import rag_prompt
13
-
14
- # Set API Keys
15
- os.environ["GROQ_API_KEY"] = st.secrets.get("GROQ_API_KEY", "")
16
-
17
- # Load LLM models
18
- llm_judge = ChatGroq(model="deepseek-r1-distill-llama-70b")
19
- rag_llm = ChatGroq(model="mixtral-8x7b-32768")
20
-
21
- llm_judge.verbose = True
22
- rag_llm.verbose = True
23
-
24
- # Clear ChromaDB cache to fix tenant issue
25
- chromadb.api.client.SharedSystemClient.clear_system_cache()
26
-
27
- st.title("Blah - 1")
28
-
29
- # **Initialize session state variables**
30
- if "pdf_path" not in st.session_state:
31
- st.session_state.pdf_path = None
32
- if "pdf_loaded" not in st.session_state:
33
- st.session_state.pdf_loaded = False
34
- if "chunked" not in st.session_state:
35
- st.session_state.chunked = False
36
- if "vector_created" not in st.session_state:
37
- st.session_state.vector_created = False
38
- if "vector_store_path" not in st.session_state:
39
- st.session_state.vector_store_path = "./chroma_langchain_db"
40
- if "vector_store" not in st.session_state:
41
- st.session_state.vector_store = None
42
- if "documents" not in st.session_state:
43
- st.session_state.documents = None
44
-
45
- # Step 1: Choose PDF Source
46
- pdf_source = st.radio("Upload or provide a link to a PDF:", ["Upload a PDF file", "Enter a PDF URL"], index=0, horizontal=True)
47
-
48
- if pdf_source == "Upload a PDF file":
49
- uploaded_file = st.file_uploader("Upload your PDF file", type="pdf")
50
- if uploaded_file:
51
- st.session_state.pdf_path = "temp.pdf"
52
- with open(st.session_state.pdf_path, "wb") as f:
53
- f.write(uploaded_file.getbuffer())
54
- st.session_state.pdf_loaded = False
55
- st.session_state.chunked = False
56
- st.session_state.vector_created = False
57
-
58
- elif pdf_source == "Enter a PDF URL":
59
- pdf_url = st.text_input("Enter PDF URL:")
60
- if pdf_url and not st.session_state.get("pdf_loaded", False):
61
- with st.spinner("Downloading PDF..."):
62
- try:
63
- response = requests.get(pdf_url)
64
- if response.status_code == 200:
65
- st.session_state.pdf_path = "temp.pdf"
66
- with open(st.session_state.pdf_path, "wb") as f:
67
- f.write(response.content)
68
- st.session_state.pdf_loaded = False
69
- st.session_state.chunked = False
70
- st.session_state.vector_created = False
71
- st.success("βœ… PDF Downloaded Successfully!")
72
- else:
73
- st.error("❌ Failed to download PDF. Check the URL.")
74
- except Exception as e:
75
- st.error(f"Error downloading PDF: {e}")
76
-
77
- # Step 2: Process PDF
78
- if st.session_state.pdf_path and not st.session_state.get("pdf_loaded", False):
79
- with st.spinner("Loading and processing PDF..."):
80
- loader = PDFPlumberLoader(st.session_state.pdf_path)
81
- docs = loader.load()
82
- st.session_state.documents = docs
83
- st.session_state.pdf_loaded = True # βœ… Prevent re-loading
84
- st.success(f"βœ… **PDF Loaded!** Total Pages: {len(docs)}")
85
-
86
- # Step 3: Chunking
87
- if st.session_state.get("pdf_loaded", False) and not st.session_state.get("chunked", False):
88
- with st.spinner("Chunking the document..."):
89
- model_name = "nomic-ai/modernbert-embed-base"
90
- embedding_model = HuggingFaceEmbeddings(model_name=model_name, model_kwargs={'device': 'cpu'}, encode_kwargs={'normalize_embeddings': False})
91
- text_splitter = SemanticChunker(embedding_model)
92
- documents = text_splitter.split_documents(st.session_state.documents)
93
- st.session_state.documents = documents # βœ… Store chunked docs
94
- st.session_state.chunked = True # βœ… Prevent re-chunking
95
- st.success(f"βœ… **Document Chunked!** Total Chunks: {len(documents)}")
96
-
97
- # Step 4: Setup Vectorstore
98
- if st.session_state.get("chunked", False) and not st.session_state.get("vector_created", False):
99
- with st.spinner("Creating vector store..."):
100
- embedding_model = HuggingFaceEmbeddings(model_name="nomic-ai/modernbert-embed-base", model_kwargs={'device': 'cpu'}, encode_kwargs={'normalize_embeddings': False})
101
-
102
- vector_store = Chroma(
103
- collection_name="deepseek_collection",
104
- collection_metadata={"hnsw:space": "cosine"},
105
- embedding_function=embedding_model,
106
- persist_directory=st.session_state.vector_store_path
107
- )
108
- vector_store.add_documents(st.session_state.documents)
109
- num_documents = len(vector_store.get()["documents"])
110
- st.session_state.vector_store = vector_store
111
- st.session_state.vector_created = True # βœ… Prevent re-creating vector store
112
- st.success(f"βœ… **Vector Store Created!** Total documents stored: {num_documents}")
113
-
114
- # Step 5: Query Input
115
- if st.session_state.get("vector_created", False) and st.session_state.get("vector_store", None):
116
- query = st.text_input("πŸ” Enter a Query:")
117
-
118
- if query and st.session_state.get("vector_created", False):
119
- with st.spinner("Retrieving relevant contexts..."):
120
- retriever = st.session_state.vector_store.as_retriever(search_type="similarity", search_kwargs={"k": 5})
121
- contexts = retriever.invoke(query)
122
- context_texts = [doc.page_content for doc in contexts]
123
-
124
- st.success(f"βœ… **Retrieved {len(context_texts)} Contexts!**")
125
- for i, text in enumerate(context_texts, 1):
126
- st.write(f"**Context {i}:** {text[:500]}...")
127
-
128
- # **Step 6: Generate Final Response**
129
- with st.spinner("Generating the final answer..."):
130
- final_prompt = PromptTemplate(input_variables=["query", "context"], template=rag_prompt)
131
- response_chain = LLMChain(llm=rag_llm, prompt=final_prompt, output_key="final_response")
132
- final_response = response_chain.invoke({"query": query, "context": context_texts})
133
-
134
- st.subheader("πŸŸ₯ RAG Final Response")
135
- st.success(final_response['final_response'])