-
Notifications
You must be signed in to change notification settings - Fork 0
Expand file tree
/
Copy pathcode.py
More file actions
67 lines (56 loc) · 2.17 KB
/
Copy pathcode.py
File metadata and controls
67 lines (56 loc) · 2.17 KB
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
import os
from typing import List, Mapping, Any
from langchain.llms.base import LLM
from langchain.callbacks.manager import CallbackManagerForLLMRun
from langchain.chains import RetrievalQA
from langchain_community.vectorstores import Chroma
from langchain_community.embeddings import SentenceTransformerEmbeddings
from langchain_community.document_loaders import TextLoader
from langchain.text_splitter import RecursiveCharacterTextSplitter
# from langchain_community.chat_models import ChatGroq
from langchain_groq import ChatGroq
# Import ChatGroq
# from langchain_groq import ChatGroq
# 1. Load Documents
with open("sample_knowledge.txt", "w") as f:
f.write("""
Faisalabad is a major industrial city in Pakistan.
It is located in the Punjab province.
The city is known for its textile industry and is often called the "Manchester of Pakistan".
The clock tower in Faisalabad is a famous landmark.
Agriculture also plays a significant role in the economy of the surrounding region.
The current time in Faisalabad is Thursday, April 17, 2025 at 11:09 PM PKT.
""")
loader = TextLoader("sample_knowledge.txt")
documents = loader.load()
# 2. Split Text into Chunks
text_splitter = RecursiveCharacterTextSplitter(chunk_size=100, chunk_overlap=20)
chunks = text_splitter.split_documents(documents)
# 3. Create Embeddings
embeddings = SentenceTransformerEmbeddings(model_name="all-mpnet-base-v2")
# 4. Create Vector Store
db = Chroma.from_documents(chunks, embeddings)
retriever = db.as_retriever(search_kwargs={"k": 3})
# 5. Initialize Groq Chat Model
llm = ChatGroq(
model="llama-3.3-70b-versatile",
temperature=0.5,
max_tokens=1024,
api_key=os.environ["GROQ_API_KEY"], # Use the environment variable
)
# 6. Create Retrieval QA Chain
qa = RetrievalQA.from_chain_type(
llm=llm,
chain_type="stuff",
retriever=retriever,
return_source_documents=True
)
# 7. Query the RAG Application
query = "this city is manchester?"
result = qa({"query": query})
print("Query:", result["query"])
print("Answer:", result["result"])
if "source_documents" in result:
print("\nSource Documents:")
for doc in result["source_documents"]:
print(doc.page_content)