|
import streamlit as st |
|
import os |
|
from langchain.callbacks.streaming_stdout import StreamingStdOutCallbackHandler |
|
from langchain.schema import StrOutputParser |
|
|
|
from datetime import datetime, timezone, timedelta |
|
|
|
from custom_llm import CustomLLM, custom_chain_with_history |
|
|
|
from typing import Optional |
|
|
|
from langchain.prompts import ChatPromptTemplate, MessagesPlaceholder |
|
from langchain_core.chat_history import BaseChatMessageHistory |
|
from langchain.memory import ConversationBufferMemory |
|
|
|
import urllib.parse as up |
|
|
|
|
|
|
|
|
|
API_TOKEN = os.getenv('HF_INFER_API') |
|
|
|
@st.cache_resource |
|
def get_llm_chain(): |
|
return custom_chain_with_history( |
|
|
|
llm=CustomLLM(repo_id="google/gemma-7b", model_type='text-generation', api_token=API_TOKEN, stop=["\n<|","<|"], temperature=0.001), |
|
|
|
memory=st.session_state.memory |
|
) |
|
|
|
|
|
if 'memory' not in st.session_state: |
|
st.session_state['memory'] = ConversationBufferMemory(return_messages=True) |
|
|
|
|
|
|
|
|
|
st.session_state.memory.chat_memory.add_ai_message("Hello, I'm AI medical consultant. How can I help you today?") |
|
|
|
|
|
if 'chain' not in st.session_state: |
|
|
|
|
|
|
|
|
|
|
|
|
|
st.session_state['chain'] = get_llm_chain() |
|
|
|
|
|
|
|
st.title("AI Medical Consultation") |
|
st.subheader("") |
|
|
|
|
|
if "messages" not in st.session_state: |
|
st.session_state.messages = [{"role":"assistant", "content":"Hello, I'm AI medical consultant. How can I help you today?"}] |
|
|
|
|
|
for message in st.session_state.messages: |
|
with st.chat_message(message["role"]): |
|
st.markdown(message["content"]) |
|
|
|
|
|
if prompt := st.chat_input("Ask me anything.."): |
|
|
|
st.chat_message("User").markdown(prompt) |
|
|
|
st.session_state.messages.append({"role": "User", "content": prompt}) |
|
|
|
response = st.session_state.chain.invoke({"question":prompt, "memory":st.session_state.memory}).split("\n<|")[0] |
|
|
|
|
|
with st.chat_message("assistant"): |
|
st.markdown(response) |
|
|
|
|
|
|
|
st.session_state.memory.save_context({"question":prompt}, {"output":response}) |
|
st.session_state.memory.chat_memory.messages = st.session_state.memory.chat_memory.messages[-15:] |
|
|
|
st.session_state.messages.append({"role": "assistant", "content": response}) |