Update app.py

app.py

@@ -1,86 +1,109 @@
+import os
+import torch
 import streamlit as st
-from huggingface_hub import InferenceClient
-
-# Initialize the InferenceClient
-client = InferenceClient("HuggingFaceH4/zephyr-7b-beta")
+from transformers import AutoModelForCausalLM, AutoTokenizer, pipeline
+from huggingface_hub import login
 
 # Streamlit app configuration
-st.set_page_config(page_title="Health Care ChatBot")
-st.title("Health Care ChatBot")
-
-# Initialize session state for messages if not present
+st.set_page_config(page_title="Medical Chatbot", layout="wide")
+st.title("Medical Chatbot")
+
+# Get the token from the environment variable
+hf_token = os.getenv("HF_TOKEN")
+if hf_token is None:
+    raise ValueError("Hugging Face token not found. Please set the HF_TOKEN environment variable.")
+
+# Authenticate with Hugging Face
+login(hf_token)
+
+# Set the random seed for reproducibility
+torch.manual_seed(0)
+
+# Supported models
+model_links = {
+    "Phi-3-mini-128k-instruct": "microsoft/Phi-3-mini-128k-instruct",
+    "Mistral-7B": "mistralai/Mistral-7B-Instruct-v0.2",
+    "Zephyr-7B": "HuggingFaceH4/zephyr-7b-beta"
+}
+
+model_info = {
+    "Phi-3-mini-128k-instruct": {
+        'description': """Phi-3-mini-128k-instruct is a large language model from Microsoft for health-related interactions. 
+        It has been optimized for instruct-based queries.""",
+        'logo': 'https://upload.wikimedia.org/wikipedia/commons/thumb/a/a6/Microsoft_logo_%282012%29.svg/200px-Microsoft_logo_%282012%29.svg.png'
+    },
+    "Mistral-7B": {
+        'description': """Mistral 7B is a large language model from Mistral AI optimized for Q&A tasks.""",
+        'logo': 'https://mistral.ai/images/logo_hubc88c4ece131b91c7cb753f40e9e1cc5_2589_256x0_resize_q97_h2_lanczos_3.webp'
+    },
+    "Zephyr-7B": {
+        'description': """Zephyr 7B is a Huggingface model, fine-tuned for helpful and instructive interactions.""",
+        'logo': 'https://huggingface.co/HuggingFaceH4/zephyr-7b-gemma-v0.1/resolve/main/thumbnail.png'
+    }
+}
+
+# Sidebar for model selection and parameters
+selected_model = st.sidebar.selectbox("Select Model", model_links.keys())
+st.sidebar.write(f"You're now chatting with **{selected_model}**")
+st.sidebar.markdown(model_info[selected_model]['description'])
+st.sidebar.image(model_info[selected_model]['logo'])
+
+# Temperature slider
+temperature = st.sidebar.slider('Temperature', 0.1, 1.0, 0.7)
+
+# Reset conversation button
+def reset_conversation():
+    st.session_state.messages = []
+    st.session_state.model = selected_model
+
+st.sidebar.button('Reset Chat', on_click=reset_conversation)
+
+# Load model and tokenizer only if it's not already loaded
+if 'model' not in st.session_state or st.session_state.model != selected_model:
+    model = AutoModelForCausalLM.from_pretrained(
+        model_links[selected_model], 
+        device_map="auto", 
+        torch_dtype="auto", 
+        trust_remote_code=True
+    )
+    tokenizer = AutoTokenizer.from_pretrained(model_links[selected_model])
+    
+    # Initialize the text generation pipeline
+    pipe = pipeline(
+        "text-generation",
+        model=model,
+        tokenizer=tokenizer
+    )
+    
+    st.session_state.model = selected_model
+    st.session_state.pipe = pipe
+
+# Initialize chat messages
 if 'messages' not in st.session_state:
     st.session_state.messages = [
-        {"role": "system", "content": "You are a knowledgeable and empathetic medical assistant providing accurate and compassionate health advice based on user input."}
+        {"role": "system", "content": "You are a medical chatbot. You should only respond to health questions!"}
     ]
 
-def respond(message, history, max_tokens, temperature, top_p):
-    # Prepare the list of messages for the chat completion
-    messages = [{"role": "system", "content": st.session_state.messages[0]["content"]}]
-
-    for val in history:
-        if val["role"] == "user":
-            messages.append({"role": "user", "content": val["content"]})
-        elif val["role"] == "assistant":
-            messages.append({"role": "assistant", "content": val["content"]})
-
-    messages.append({"role": "user", "content": message})
-
-    # Generate response
-    response = ""
-    response_container = st.empty()  # Placeholder to update the response text dynamically
-
-    for message in client.chat_completion(
-        messages,
-        max_tokens=max_tokens,
-        stream=True,
-        temperature=temperature,
-        top_p=top_p,
-    ):
-        token = message.choices[0].delta.content
-        response += token
-        response_container.text(response)  # Stream the response
-
-    return response
-
-# Sidebar for parameters
-with st.sidebar:
-    max_tokens = st.slider("Max new tokens", 1, 2048, 512)
-    temperature = st.slider("Temperature", 0.1, 4.0, 0.7)
-    top_p = st.slider("Top-p (nucleus sampling)", 0.1, 1.0, 0.95)
-
-# Display chat messages from history
+# Display chat history
 for message in st.session_state.messages:
-    if message["role"] == "user":
-        # User message on the right
-        col1, col2 = st.columns([1, 4])
-        with col2:
-            with st.chat_message("user"):
-                st.write(message["content"])
-        with col1:
-            st.write("")  # Empty space on the left for alignment
-        
-    elif message["role"] == "assistant":
-        # Assistant message on the left
-        col1, col2 = st.columns([4, 1])
-        with col1:
-            with st.chat_message("assistant"):
-                st.write(message["content"])
-        with col2:
-            st.write("")  # Empty space on the right for alignment
-
-# Keep user input box at the bottom
-st.divider()  # Optional, to visually separate chat history from input box
-user_input = st.text_input("You:", key="user_message", placeholder="Type your message here...")
-
-if user_input:
-    # Append user message to the chat history
+    with st.chat_message(message["role"]):
+        st.markdown(message["content"])
+
+# Function to generate responses
+def generate_response(messages):
+    messages_str = "\n".join([f"{msg['role']}: {msg['content']}" for msg in messages if msg['role'] != 'system'])
+    output = st.session_state.pipe(messages_str, max_new_tokens=150, temperature=temperature, return_full_text=False)
+    return output[0]['generated_text']
+
+# User input
+if user_input := st.chat_input("Ask a health question..."):
+    # Display user message
+    with st.chat_message("user"):
+        st.markdown(user_input)
     st.session_state.messages.append({"role": "user", "content": user_input})
 
-    # Generate assistant response
-    response = respond(user_input, st.session_state.messages, max_tokens, temperature, top_p)
+    # Generate and display assistant response
+    response = generate_response(st.session_state.messages)
+    with st.chat_message("assistant"):
+        st.markdown(response)
     st.session_state.messages.append({"role": "assistant", "content": response})
-
-    # Refresh to display new messages
-    st.experimental_rerun()
-