Spaces:
Sleeping
Sleeping
File size: 5,364 Bytes
7388333 |
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 |
import streamlit as st
import time
from autism_chatbot import *
class StreamHandler:
def __init__(self, placeholder):
self.text = ""
self.text_container = placeholder
def append_text(self, text: str) -> None:
self.text += text
self.text_container.markdown(self.text)
class StreamingGroqLLM(GroqLLM):
stream_handler: Any = Field(None, description="Stream handler for real-time output")
def _call(self, prompt: str, stop: Optional[List[str]] = None, **kwargs: Any) -> str:
completion = self.client.chat.completions.create(
messages=[{"role": "user", "content": prompt}],
model=self.model_name,
stream=True,
**kwargs
)
collected_chunks = []
collected_messages = []
for chunk in completion:
chunk_message = chunk.choices[0].delta.content
if chunk_message is not None:
collected_chunks.append(chunk_message)
collected_messages.append(chunk_message)
if self.stream_handler:
self.stream_handler.append_text(chunk_message)
time.sleep(0.05)
return ''.join(collected_messages)
class StreamingAutismResearchBot(AutismResearchBot):
def __init__(self, groq_api_key: str, stream_handler: StreamHandler, index_path: str = "index.faiss"):
self.llm = StreamingGroqLLM(
groq_api_key=groq_api_key,
model_name="llama-3.3-70b-versatile",
stream_handler=stream_handler
)
self.embeddings = HuggingFaceEmbeddings(
model_name="pritamdeka/S-PubMedBert-MS-MARCO",
model_kwargs={'device': 'cpu'}
)
self.db = FAISS.load_local("./", self.embeddings, allow_dangerous_deserialization=True)
self.memory = ConversationBufferMemory(
memory_key="chat_history",
return_messages=True,
output_key="answer"
)
self.qa_chain = self._create_qa_chain()
def main():
# Page configuration
st.set_page_config(
page_title="Autism Research Assistant",
page_icon="🧩",
layout="wide"
)
# Add custom CSS
st.markdown("""
<style>
.stApp {
max-width: 1200px;
margin: 0 auto;
}
.stMarkdown {
font-size: 16px;
}
.chat-message {
padding: 1rem;
border-radius: 0.5rem;
margin-bottom: 1rem;
}
.timestamp {
font-size: 0.8em;
color: #666;
}
</style>
""", unsafe_allow_html=True)
# Header
st.title("🧩 Autism Research Assistant")
st.markdown("""
Welcome to your AI-powered autism research assistant. I'm here to provide evidence-based
assessments and therapy recommendations based on scientific research.
""")
# Initialize session state
if 'messages' not in st.session_state:
st.session_state.messages = [
{"role": "assistant", "content": "Hello! I'm your autism research assistant. How can I help you today?"}
]
# Initialize bot
if 'bot' not in st.session_state:
st.session_state.stream_container = None
st.session_state.bot = None
# Display chat messages
for message in st.session_state.messages:
with st.chat_message(message["role"]):
st.write(f"{message['content']}")
st.caption(f"{time.strftime('%I:%M %p')}")
# Chat input
if prompt := st.chat_input("Type your message here..."):
# Display user message
with st.chat_message("user"):
st.write(prompt)
st.caption(f"{time.strftime('%I:%M %p')}")
# Add to session state
st.session_state.messages.append({"role": "user", "content": prompt})
# Create a new chat message container for the assistant's response
assistant_message = st.chat_message("assistant")
with assistant_message:
# Create a placeholder for the streaming text
stream_placeholder = st.empty()
# Initialize the bot with the new stream handler if not already initialized
if st.session_state.bot is None:
stream_handler = StreamHandler(stream_placeholder)
st.session_state.bot = StreamingAutismResearchBot(
groq_api_key= os.environ.get("GROQ_API_KEY"),
stream_handler=stream_handler,
)
else:
# Update the stream handler with the new placeholder
st.session_state.bot.llm.stream_handler.text = ""
st.session_state.bot.llm.stream_handler.text_container = stream_placeholder
# Generate response
response = st.session_state.bot.answer_question(prompt)
# Clear the streaming placeholder and display the final message
stream_placeholder.empty()
st.write(response['answer'])
st.caption(f"{time.strftime('%I:%M %p')}")
# Add bot response to session state
st.session_state.messages.append({"role": "assistant", "content": response['answer']})
if __name__ == "__main__":
main() |