updating app.py

app.py

@@ -5,6 +5,19 @@ import re
 import requests
 import uuid
 import time
+from transformers import AutoModelForCausalLM, AutoTokenizer, pipeline
+
+# Load the Hugging Face model
+model_name = "ethicsadvisorproject/Llama-2-7b-ethical-chat-finetune"
+tokenizer = AutoTokenizer.from_pretrained(model_name, cache_dir="/tmp")
+model = AutoModelForCausalLM.from_pretrained(
+    model_name,
+    torch_dtype=torch.float16,
+    device_map="auto",
+    offload_folder="/tmp"
+)
+
+pipe = pipeline("text-generation", model=model, tokenizer=tokenizer, max_length=200)
 
 DB_DIR = 'user_data'  # Directory to store individual user data
 os.makedirs(DB_DIR, exist_ok=True)  # Ensure the directory exists
@@ -89,34 +102,43 @@ def main():
         with st.chat_message(message["role"]):
             st.markdown(message["content"])
 
-    # User input
-    if prompt := st.chat_input("What is up?"):
-        st.session_state.messages.append({"role": "user", "content": prompt})
-        with st.chat_message("user"):
-            st.markdown(prompt)
+    # # User input
+    # # if prompt := st.chat_input("What is up?"):
+    # #     st.session_state.messages.append({"role": "user", "content": prompt})
+    # #     with st.chat_message("user"):
+    # #         st.markdown(prompt)
+
+    # #     # Send request to the endpoint
+    # #     headers = {'ngrok-skip-browser-warning': 'true'}
+    # #     data = {'messages': st.session_state.messages[-1]['content']}
 
-        # Send request to the endpoint
-        headers = {'ngrok-skip-browser-warning': 'true'}
-        data = {'messages': st.session_state.messages[-1]['content']}
+    # #     try:
+    # #         response = requests.post(endpoint_url, json=data, headers=headers)
+    # #         response.raise_for_status()  # Raise exception for HTTP errors
+    # #         response_data = response.json()
+    # #         response_text = response_data.get('response_text', '')
 
-        try:
-            response = requests.post(endpoint_url, json=data, headers=headers)
-            response.raise_for_status()  # Raise exception for HTTP errors
-            response_data = response.json()
-            response_text = response_data.get('response_text', '')
+    #         # Clean response text
+    #         message = re.sub(r'<s>\[INST\].*?\[/INST\]', '', response_text).strip()
 
-            # Clean response text
-            message = re.sub(r'<s>\[INST\].*?\[/INST\]', '', response_text).strip()
+    #         with st.chat_message("assistant"):
+    #             st.markdown(message)
 
-            with st.chat_message("assistant"):
-                st.markdown(message)
+    #         st.session_state.messages.append({"role": "assistant", "content": message})
+
+    #     except requests.exceptions.RequestException as e:
+    #         st.error(f"Error communicating with the endpoint: {e}")
+    #     except KeyError:
+    #         st.error(f"Unexpected response format. Missing 'response_text' key. Received: {response.text}")
+
+    if prompt := st.chat_input("What is up?"):
+    response = pipe(f"<s>[INST] {prompt} [/INST]")
+    response_text = response[0]["generated_text"].replace("<s>[INST]", "").replace("[/INST]", "").strip()
 
-            st.session_state.messages.append({"role": "assistant", "content": message})
+    with st.chat_message("assistant"):
+        st.markdown(response_text)
 
-        except requests.exceptions.RequestException as e:
-            st.error(f"Error communicating with the endpoint: {e}")
-        except KeyError:
-            st.error(f"Unexpected response format. Missing 'response_text' key. Received: {response.text}")
+    st.session_state.messages.append({"role": "assistant", "content": response_text})
 
         # Save updated chat history
         user_data["chat_history"] = st.session_state.messages