Spaces:
Sleeping
Sleeping
File size: 4,043 Bytes
74bd715 |
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 |
import os
import zipfile
import chromadb
import gradio as gr
from langchain_core.prompts import ChatPromptTemplate
from langchain_core.runnables import RunnableLambda, 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.")
# Initialize ChromaDB client
chroma_client = chromadb.PersistentClient(path="./chroma_store")
# Vector store and retriever
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 = vectorstore.as_retriever(search_kwargs={"k": 3, "filter": {"site": "imageonline"}})
# Retrieval logic
def retrieve_with_metadata(query, k=5):
docs = retriever.get_relevant_documents(query)
if not docs:
return {"context": "No relevant context found.", "references": []}
top_doc = docs[0]
return {
"context": top_doc.page_content,
"references": [{
"section": top_doc.metadata.get("section", "Unknown"),
"source": top_doc.metadata.get("source", "Unknown")
}]
}
# LLM setup
llm = ChatTogether(
model="meta-llama/Llama-3-8b-chat-hf",
temperature=0.3,
max_tokens=1024,
top_p=0.7,
together_api_key="a36246d65d8290f43667350b364c5b6bb8562eb50a4b947eec5bd7e79f2dffc6"
)
# Prompt template
prompt = ChatPromptTemplate.from_template("""
You are an expert assistant for ImageOnline Web Solutions.
Answer the user's query based ONLY on the following context:
{context}
Query: {question}
""")
rag_chain = (
{
"context": lambda x: retrieve_with_metadata(x)["context"],
"question": RunnablePassthrough()
}
| prompt
| llm
| StrOutputParser()
)
def get_references(query):
return retrieve_with_metadata(query)["references"]
# Gradio UI
def chat_interface(message, history):
history = history or []
history.append(("🧑 You: " + message, "⏳ Generating response..."))
try:
answer = rag_chain.invoke(message)
references = get_references(message)
if references:
ref = references[0]
ref_string = f"\n\n📚 **Reference:**\nSection: {ref['section']}\nURL: {ref['source']}"
else:
ref_string = "\n\n📚 **Reference:**\n_None available_"
full_response = answer + ref_string
history[-1] = ("🧑 You: " + message, "🤖 Bot: " + full_response)
except Exception as e:
history[-1] = ("🧑 You: " + message, f"🤖 Bot: ⚠️ {str(e)}")
return history, history
def launch_gradio():
with gr.Blocks() as demo:
gr.Markdown("# 💬 ImageOnline RAG Chatbot")
gr.Markdown("Ask about Website Designing, App Development, SEO, Hosting, etc.")
chatbot = gr.Chatbot()
state = gr.State([])
with gr.Row():
msg = gr.Textbox(placeholder="Ask your question here...", show_label=False, scale=8)
send_btn = gr.Button("📨 Send", scale=1)
msg.submit(chat_interface, inputs=[msg, state], outputs=[chatbot, state])
send_btn.click(chat_interface, inputs=[msg, state], outputs=[chatbot, state])
with gr.Row():
clear_btn = gr.Button("🧹 Clear Chat")
clear_btn.click(fn=lambda: ([], []), outputs=[chatbot, state])
return demo
demo = launch_gradio()
demo.launch() |