Spaces:
No application file
No application file
File size: 4,254 Bytes
156bb4a | 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 | import os
from dotenv import load_dotenv
import gradio as gr
from langchain_community.embeddings import HuggingFaceBgeEmbeddings
from langchain_community.document_loaders import PyPDFLoader, DirectoryLoader
from langchain.text_splitter import RecursiveCharacterTextSplitter
from langchain_community.vectorstores import Chroma
from langchain.chains import RetrievalQA
from langchain.prompts import PromptTemplate
from langchain_groq import ChatGroq
# Load API keys securely
load_dotenv()
GROQ_API_KEY = os.getenv("GROQ_API_KEY")
HUGGINGFACEHUB_API_TOKEN = os.getenv("HUGGINGFACEHUB_API_TOKEN")
if not GROQ_API_KEY or not HUGGINGFACEHUB_API_TOKEN:
raise ValueError("❌ Missing API keys! Make sure to set them in a .env file.")
def initialize_llm():
"""Initialize the LLaMA model using Groq API."""
return ChatGroq(
temperature=0,
model_name='llama-3.3-70b-versatile',
groq_api_key=GROQ_API_KEY
)
def create_vector_db():
"""Load PDFs, process text, and create a Chroma vector database."""
loader = DirectoryLoader('./sample_data', glob='*.pdf', loader_cls=PyPDFLoader)
# Process texts and create embeddings here
documents = loader.load()
text_splitter = RecursiveCharacterTextSplitter(chunk_size=500, chunk_overlap=50)
texts = text_splitter.split_documents(documents)
embeddings = HuggingFaceBgeEmbeddings(
model_name="sentence-transformers/all-MiniLM-L6-v2"
)
db = Chroma.from_documents(texts, embeddings, persist_directory="chroma_db")
db.persist()
print("✅ ChromaDB created successfully!")
return db
def setup_qa_chain(db, llm):
"""Set up the RetrievalQA chain."""
retriever = db.as_retriever()
prompt_template = """You are a mental health expert. Use the following information to answer the user's question.
If you don't know the answer, just say that you don't know, don't try to make up an answer.
{context}
user={question}
chatbot:"""
PROMPT = PromptTemplate(template=prompt_template, input_variables=["context", "question"])
qa_chain = RetrievalQA.from_chain_type(
llm=llm,
chain_type="stuff",
retriever=retriever,
chain_type_kwargs={"prompt": PROMPT}
)
return qa_chain
# Initialize components
print("🔄 Initializing CalmMateAI...")
llm = initialize_llm()
db_path = "chroma_db"
if os.path.exists(db_path) and os.listdir(db_path):
print("🟢 Loading existing ChromaDB...")
embedding = HuggingFaceBgeEmbeddings(model_name="sentence-transformers/all-MiniLM-L6-v2")
vector_db = Chroma(persist_directory=db_path, embedding_function=embedding)
else:
print("🚀 Creating a new ChromaDB...")
vector_db = create_vector_db()
qa_chain = setup_qa_chain(vector_db, llm)
# Gradio Chat Function
def chat_response(user_input, history=[]):
if not user_input.strip():
return history + [("You", user_input), ("CalmMateAI", "Please enter a valid question. 🙏")], ""
# Append user input to history
history.append(("You", user_input))
# Construct the context from history
context = " ".join([f"{role}: {text}" for role, text in history])
try:
# Generate the response using the context
response = qa_chain.invoke({"query": user_input, "context": context})["result"]
except Exception as e:
response = f"⚠️ Error: {str(e)}"
# Append response to history
history.append(("CalmMateAI", response))
return history, "" # Clears input field after sending
# Gradio UI
theme = gr.themes.Soft(primary_hue="blue", secondary_hue="gray")
with gr.Blocks(theme=theme) as app:
gr.Markdown("# 🌿 CalmMateAI: Your Mental Health Companion")
gr.Markdown("**A safe space for mental well-being. Ask anything, and I'll help!**")
with gr.Row():
chatbot = gr.Chatbot(label="CalmMateAI Chat", height=400)
user_input = gr.Textbox(placeholder="Type your question here...", show_label=False)
send_button = gr.Button("Send", variant="primary")
send_button.click(chat_response, inputs=[user_input, chatbot], outputs=[chatbot, user_input])
user_input.submit(chat_response, inputs=[user_input, chatbot], outputs=[chatbot, user_input])
app.launch(debug=True)
|