Spaces:
Running
Running
import streamlit as st | |
import os | |
import time | |
from services.llm import initialize_llm, initialize_embeddings | |
from services.vector_store import create_vector_store, retrive_vector_store, generate_prompt | |
from services.pdf_processing import load_and_split_pdf | |
from utils.helpers import extract_thoughts, response_generator | |
import subprocess | |
# try: | |
# print("π Checking and starting Ollama...") | |
# subprocess.run(["bash", "install_ollama.sh"], check=True) | |
# print("β Ollama is running!") | |
# except subprocess.CalledProcessError as e: | |
# print(f"β Error: {e}") | |
# Custom CSS for chat styling | |
CHAT_CSS = """ | |
<style> | |
.user-message { | |
text-align: right; | |
background-color: #3c8ce7; | |
color: white; | |
padding: 10px; | |
border-radius: 10px; | |
margin-bottom: 10px; | |
display: inline-block; | |
width: fit-content; | |
max-width: 70%; | |
margin-left: auto; | |
box-shadow: 0px 4px 6px rgba(0, 0, 0, 0.1); | |
} | |
.assistant-message { | |
text-align: left; | |
background-color: #d16ba5; | |
color: white; | |
padding: 10px; | |
border-radius: 10px; | |
margin-bottom: 10px; | |
display: inline-block; | |
width: fit-content; | |
max-width: 70%; | |
margin-right: auto; | |
box-shadow: 0px 4px 6px rgba(0, 0, 0, 0.1); | |
} | |
</style> | |
""" | |
# Streamlit UI Setup | |
st.set_page_config(page_title="DocChatAI", layout="wide") | |
st.title("π DocChatAI | Chat Using Documents") | |
# Expandable Disclaimer Section | |
with st.expander("β οΈ Disclaimer (Click to expand)"): | |
st.markdown(""" | |
- This AI chatbot provides **informational responses only** and should not be used as **legal, medical, or financial advice**. | |
- The accuracy of responses depends on the provided **context and training data**. | |
- **Use at your own discretion** and always verify important information from reliable sources. | |
""") | |
# Sidebar | |
st.sidebar.title("DocChatAI") | |
st.sidebar.subheader("Chat using PDF Document") | |
st.sidebar.write("---") | |
# Model Selection | |
selected_model = st.sidebar.radio("Choose Model", ["deepseek-r1:1.5b"]) | |
st.sidebar.write("---") | |
# Hyperparameters | |
temperature = st.sidebar.slider("Temperature", 0.0, 1.0, 0.7, 0.1) | |
top_p = st.sidebar.slider("Top-p (Nucleus Sampling)", 0.0, 1.0, 0.9, 0.05) | |
max_tokens = st.sidebar.number_input("Max Tokens", 10, 2048, 1024, 10) | |
st.sidebar.write("---") | |
# File Upload | |
uploaded_file = st.sidebar.file_uploader("π Upload a PDF", type=["pdf"]) | |
st.sidebar.write("---") | |
# About Section | |
st.sidebar.write("π **About Me**") | |
st.sidebar.write("π€ **Name:** Deepak Yadav") | |
st.sidebar.write("π‘ **Bio:** Passionate about AI and Machine Learning.") | |
st.sidebar.markdown("[GitHub](https://github.com/deepak7376) | [LinkedIn](https://www.linkedin.com/in/dky7376/)") | |
st.sidebar.write("---") | |
# Initialize LLM | |
llm = initialize_llm(selected_model, temperature, top_p, max_tokens) | |
embeddings = initialize_embeddings() | |
# Document Handling | |
retriever = None | |
if uploaded_file: | |
os.makedirs("docs", exist_ok=True) | |
filepath = os.path.join("docs", uploaded_file.name) | |
with open(filepath, "wb") as f: | |
f.write(uploaded_file.read()) | |
# Load and process PDF | |
splits = load_and_split_pdf(filepath) | |
vectorstore = create_vector_store(splits, embeddings) | |
retriever = vectorstore.as_retriever(search_type="similarity", search_kwargs={"k": 3}) | |
# Apply custom CSS | |
st.markdown(CHAT_CSS, unsafe_allow_html=True) | |
# Initialize chat history | |
if "messages" not in st.session_state: | |
st.session_state.messages = [] | |
# Display previous messages | |
for message in st.session_state.messages: | |
if message['thinking_part']: | |
with st.expander("π Thought Process"): | |
st.markdown(message['thinking_part']) | |
with st.chat_message(message["role"]): | |
st.markdown(message["content"]) | |
# Chat Input | |
if user_input := st.chat_input("π¬ Ask something..."): | |
st.session_state.messages.append({"role": "user", "content": user_input, "thinking_part": False}) | |
with st.chat_message("user"): | |
st.markdown(user_input) | |
# Measure response time | |
start_time = time.time() | |
# Generate response | |
context = retrive_vector_store(retriever, user_input) if retriever else "No context" | |
query = generate_prompt(context=context, question=user_input) | |
# response = llm.invoke(query) | |
response = llm.create_chat_completion( | |
messages = [ | |
{ | |
"role": "user", | |
"content": f"{query}" | |
} | |
] | |
) | |
# Calculate response time | |
response_time = round(time.time() - start_time, 2) | |
# Extract thoughts and main answer | |
thinking_part, main_answer = extract_thoughts(response['choices'][0]['message']['content']) | |
# Display AI response | |
with st.chat_message("assistant"): | |
if thinking_part: | |
with st.expander("π Thought Process"): | |
st.markdown(thinking_part) | |
# **Formatted Response Display** | |
formatted_response = f""" | |
{main_answer} | |
β³ **Response Time:** {response_time} seconds | |
""" | |
st.markdown(formatted_response, unsafe_allow_html=True) | |
# Save to session history | |
st.session_state.messages.append({"role": "assistant", "content": formatted_response, "thinking_part": thinking_part}) | |