Spaces:
Sleeping
Sleeping
File size: 6,036 Bytes
a855c64 |
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 |
import os
import zipfile
import chromadb
import gradio as gr
from langchain_core.prompts import ChatPromptTemplate
from langchain_core.runnables import RunnablePassthrough
from langchain_core.output_parsers import StrOutputParser
from langchain_together import ChatTogether
from langchain_community.vectorstores import Chroma
from langchain_community.embeddings import HuggingFaceEmbeddings
# Log: Check if chroma_store exists
if not os.path.exists("chroma_store"):
print("π chroma_store folder not found. Attempting to unzip...")
try:
with zipfile.ZipFile("chroma_store.zip", "r") as zip_ref:
zip_ref.extractall("chroma_store")
print("β
Successfully extracted chroma_store.zip.")
except Exception as e:
print(f"β Failed to unzip chroma_store.zip: {e}")
else:
print("β
chroma_store folder already exists. Skipping unzip.")
# ChromaDB setup
chroma_client = chromadb.PersistentClient(path="./chroma_store")
embedding_function = HuggingFaceEmbeddings(model_name="BAAI/bge-base-en-v1.5")
vectorstore = Chroma(
client=chroma_client,
collection_name="imageonline_chunks",
embedding_function=embedding_function
)
# Retriever setup (k=5)
retriever = vectorstore.as_retriever(search_kwargs={"k": 5, "filter": {"site": "imageonline"}})
# Updated retrieval logic: return full concatenated context and top 2 references
def retrieve_with_metadata(query, k=5, max_refs=2):
docs = retriever.get_relevant_documents(query)
if not docs:
return {
"context": "No relevant context found.",
"references": []
}
# Join all documents for LLM input
context = "\n\n".join(doc.page_content for doc in docs)
# Unique references (max 2)
seen = set()
references = []
for doc in docs:
source = doc.metadata.get("source", "Unknown")
section = doc.metadata.get("section", "Unknown")
key = (section, source)
if key not in seen:
seen.add(key)
references.append({"section": section, "source": source})
if len(references) >= max_refs:
break
return {
"context": context,
"references": references
}
# LLM initialization
llm = ChatTogether(
model="meta-llama/Llama-3-8b-chat-hf",
temperature=0.3,
max_tokens=1024,
top_p=0.7,
together_api_key="a36246d65d8290f43667350b364c5b6bb8562eb50a4b947eec5bd7e79f2dffc6"
)
# Improved Prompt Template
prompt = ChatPromptTemplate.from_template("""
You are a knowledgeable assistant for ImageOnline Pvt. Ltd. .
Answer the user's query using ONLY the following context extracted from our official website.
If the answer is not clearly present in the context, say "I couldn't find the information on the site."
--------------------
{context}
--------------------
Query: {question}
""")
# RAG chain
rag_chain = (
{
"context": lambda x: retrieve_with_metadata(x)["context"],
"question": RunnablePassthrough()
}
| prompt
| llm
| StrOutputParser()
)
# References for display
def get_references(query):
return retrieve_with_metadata(query)["references"]
from datetime import datetime
import time
import gradio as gr
# Chat function
def chat_interface(message, history):
history = history or []
timestamp_user = datetime.now().strftime("%H:%M:%S")
user_msg = f"π§ **You**\n{message}\n\n<span style='font-size: 0.8em; color: gray;'>β±οΈ {timestamp_user}</span>"
bot_msg = "β³ _Bot is typing..._"
history.append((user_msg, bot_msg))
try:
time.sleep(0.5)
answer = rag_chain.invoke(message)
references = get_references(message)
if references:
ref_lines = "\n".join(f"{ref['section']} β {ref['source']}" for ref in references)
ref_string = f"\n\nπ **Reference(s):**\n{ref_lines}"
else:
ref_string = "\n\nπ **Reference(s):**\n_None available_"
full_response = answer.strip() + ref_string
timestamp_bot = datetime.now().strftime("%H:%M:%S")
bot_msg = f"π€ **Bot**\n{full_response}\n\n<span style='font-size: 0.8em; color: gray;'>β±οΈ {timestamp_bot}</span>"
history[-1] = (user_msg, bot_msg)
except Exception as e:
timestamp_bot = datetime.now().strftime("%H:%M:%S")
error_msg = f"π€ **Bot**\nβ οΈ {str(e)}\n\n<span style='font-size: 0.8em; color: gray;'>β±οΈ {timestamp_bot}</span>"
history[-1] = (user_msg, error_msg)
return history, history, ""
# Gradio Launcher
def launch_gradio():
with gr.Blocks(css="""
.gr-button {
background-color: orange !important;
color: white !important;
font-weight: bold;
border-radius: 6px !important;
border: 1px solid darkorange !important;
}
.gr-button:hover {
background-color: darkorange !important;
}
.gr-textbox textarea {
border: 2px solid orange !important;
border-radius: 6px !important;
padding: 0.75rem !important;
font-size: 1rem;
}
""") as demo:
# Header and Subtitle
gr.Markdown("# π¬ ImageOnline RAG Chatbot")
gr.Markdown("Welcome! Ask about Website Designing, Web Development, App Development, About Us, Digital Marketing etc.")
chatbot = gr.Chatbot()
state = gr.State([])
with gr.Row(equal_height=True):
msg = gr.Textbox(
placeholder="Ask your question here...",
show_label=False,
scale=9
)
send_btn = gr.Button("π Send", scale=1)
msg.submit(chat_interface, inputs=[msg, state], outputs=[chatbot, state, msg])
send_btn.click(chat_interface, inputs=[msg, state], outputs=[chatbot, state, msg])
with gr.Row():
clear_btn = gr.Button("π§Ή Clear Chat")
clear_btn.click(fn=lambda: ([], []), outputs=[chatbot, state])
return demo
# Launch the app
demo = launch_gradio()
demo.launch() |