CCockrum commited on
Commit
8b550ae
Β·
verified Β·
1 Parent(s): 062b2d3

Update app.py

Browse files
Files changed (1) hide show
  1. app.py +45 -138
app.py CHANGED
@@ -1,199 +1,106 @@
1
  import os
2
  import re
3
  import requests
4
- import torch
5
  import streamlit as st
6
  from langchain_huggingface import HuggingFaceEndpoint
7
  from langchain_core.prompts import PromptTemplate
8
  from langchain_core.output_parsers import StrOutputParser
9
  from transformers import pipeline
10
- from langdetect import detect # Ensure this package is installed
11
-
12
- # βœ… Check for GPU or Default to CPU
13
- device = "cuda" if torch.cuda.is_available() else "cpu"
14
- print(f"βœ… Using device: {device}") # Debugging info
15
 
16
  # βœ… Environment Variables
17
  HF_TOKEN = os.getenv("HF_TOKEN")
18
- if HF_TOKEN is None:
 
 
19
  raise ValueError("HF_TOKEN is not set. Please add it to your environment variables.")
20
 
21
- NASA_API_KEY = os.getenv("NASA_API_KEY")
22
- if NASA_API_KEY is None:
23
  raise ValueError("NASA_API_KEY is not set. Please add it to your environment variables.")
24
 
25
  # βœ… Set Up Streamlit
26
  st.set_page_config(page_title="HAL - NASA ChatBot", page_icon="πŸš€")
27
 
28
- # βœ… Initialize Session State Variables (Ensuring Chat History Persists)
29
  if "chat_history" not in st.session_state:
30
  st.session_state.chat_history = [{"role": "assistant", "content": "Hello! How can I assist you today?"}]
31
 
32
- # βœ… Initialize Hugging Face Model (Explicitly Set to CPU/GPU)
33
  def get_llm_hf_inference(model_id="mistralai/Mistral-7B-Instruct-v0.3", max_new_tokens=512, temperature=0.7):
34
  return HuggingFaceEndpoint(
35
  repo_id=model_id,
36
  max_new_tokens=max_new_tokens,
37
  temperature=temperature,
38
  token=HF_TOKEN,
39
- task="text-generation",
40
- device=-1 if device == "cpu" else 0 # βœ… Force CPU (-1) or GPU (0)
41
  )
42
 
43
- # βœ… NASA API Function
44
- def get_nasa_apod():
45
- url = f"https://api.nasa.gov/planetary/apod?api_key={NASA_API_KEY}"
46
- response = requests.get(url)
47
- if response.status_code == 200:
48
- data = response.json()
49
- return data.get("url", ""), data.get("title", ""), data.get("explanation", "")
50
- return "", "NASA Data Unavailable", "I couldn't fetch data from NASA right now."
51
-
52
- # βœ… Sentiment Analysis (Now Uses Explicit Device)
53
- sentiment_analyzer = pipeline(
54
- "sentiment-analysis",
55
- model="distilbert/distilbert-base-uncased-finetuned-sst-2-english",
56
- device=-1 if device == "cpu" else 0 # βœ… Force CPU (-1) or GPU (0)
57
- )
58
-
59
- def analyze_sentiment(user_text):
60
- result = sentiment_analyzer(user_text)[0]
61
- return result['label']
62
-
63
- # βœ… Intent Detection
64
- def predict_action(user_text):
65
- if "NASA" in user_text.lower() or "space" in user_text.lower():
66
- return "nasa_info"
67
- return "general_query"
68
-
69
- # βœ… Ensure English Responses
70
- def ensure_english(text):
71
- try:
72
- detected_lang = detect(text)
73
- if detected_lang != "en":
74
- return "⚠️ Sorry, I only respond in English. Can you rephrase your question?"
75
- except:
76
- return "⚠️ Language detection failed. Please ask your question again."
77
- return text
78
-
79
- # βœ… Follow-Up Question Generation
80
  def generate_follow_up(user_text):
81
- """Generates a structured follow-up question in a concise format."""
82
-
83
  prompt_text = (
84
  f"Given the user's question: '{user_text}', generate a SHORT follow-up question in the format: "
85
- "'Would you like to learn more about [related topic] or explore something else?'. "
86
- "Ensure it's concise and structured exactly as requested without extra commentary."
87
  )
88
-
89
- hf = get_llm_hf_inference(max_new_tokens=30, temperature=0.6) # πŸ”₯ Lower temp for consistency
90
  output = hf.invoke(input=prompt_text).strip()
91
 
92
- # βœ… Extract the relevant part using regex to remove unwanted symbols or truncations
93
  cleaned_output = re.sub(r"```|''|\"", "", output).strip()
94
 
95
- # βœ… Ensure output is formatted correctly
96
- if "Would you like to learn more about" not in cleaned_output:
97
- cleaned_output = "Would you like to explore another related topic or ask about something else?"
98
-
99
- return cleaned_output
100
-
101
- # βœ… Main Response Function
102
- def get_response(system_message, chat_history, user_text, max_new_tokens=512):
103
- action = predict_action(user_text)
104
-
105
- # βœ… Handle NASA-Specific Queries
106
- if action == "nasa_info":
107
- nasa_url, nasa_title, nasa_explanation = get_nasa_apod()
108
- response = f"**{nasa_title}**\n\n{nasa_explanation}"
109
- follow_up = generate_follow_up(user_text)
110
- chat_history.extend([
111
- {'role': 'user', 'content': user_text},
112
- {'role': 'assistant', 'content': response},
113
- {'role': 'assistant', 'content': follow_up}
114
- ])
115
- return response, follow_up, chat_history, nasa_url
116
 
117
- # βœ… Invoke Hugging Face Model
118
- hf = get_llm_hf_inference(max_new_tokens=max_new_tokens, temperature=0.9)
 
 
 
119
 
120
- filtered_history = "\n".join(f"{msg['role']}: {msg['content']}" for msg in chat_history)
 
121
 
 
122
  prompt = PromptTemplate.from_template(
123
- "[INST] {system_message}\n\nCurrent Conversation:\n{chat_history}\n\n"
124
  "User: {user_text}.\n [/INST]\n"
125
- "AI: Provide a detailed explanation. Use a conversational tone. "
126
- "🚨 Answer **only in English**."
127
  "\nHAL:"
128
  )
129
 
130
  chat = prompt | hf.bind(skip_prompt=True) | StrOutputParser(output_key='content')
131
- response = chat.invoke(input=dict(system_message=system_message, user_text=user_text, chat_history=filtered_history))
132
  response = response.split("HAL:")[-1].strip() if "HAL:" in response else response.strip()
133
 
134
- response = ensure_english(response)
135
-
136
- if not response:
137
- response = "I'm sorry, but I couldn't generate a response. Can you rephrase your question?"
138
-
139
  follow_up = generate_follow_up(user_text)
140
 
141
- chat_history.extend([
142
- {'role': 'user', 'content': user_text},
143
- {'role': 'assistant', 'content': response},
144
- {'role': 'assistant', 'content': follow_up}
145
- ])
146
 
147
- return response, follow_up, chat_history, None
148
 
149
- # βœ… Streamlit UI
150
  st.title("πŸš€ HAL - NASA AI Assistant")
151
 
152
- # βœ… Justify all chatbot responses
153
- st.markdown("""
154
- <style>
155
- .user-msg, .assistant-msg {
156
- padding: 10px;
157
- border-radius: 10px;
158
- margin-bottom: 5px;
159
- width: fit-content;
160
- max-width: 80%;
161
- text-align: justify;
162
- }
163
- .user-msg { background-color: #696969; color: white; }
164
- .assistant-msg { background-color: #333333; color: white; }
165
- .container { display: flex; flex-direction: column; align-items: flex-start; }
166
- @media (max-width: 600px) { .user-msg, .assistant-msg { font-size: 16px; max-width: 100%; } }
167
- </style>
168
- """, unsafe_allow_html=True)
169
-
170
- # βœ… Reset Chat Button
171
- if st.sidebar.button("Reset Chat"):
172
- st.session_state.chat_history = [{"role": "assistant", "content": "Hello! How can I assist you today?"}]
173
- st.session_state.response_ready = False
174
- st.session_state.follow_up = ""
175
 
176
- # βœ… Chat UI
177
  user_input = st.chat_input("Type your message here...")
178
 
179
  if user_input:
180
- response, follow_up, st.session_state.chat_history, image_url = get_response(
181
- system_message="You are a helpful AI assistant.",
182
- user_text=user_input,
183
- chat_history=st.session_state.chat_history
184
- )
185
-
186
- if response:
187
- st.markdown(f"<div class='assistant-msg'><strong>HAL:</strong> {response}</div>", unsafe_allow_html=True)
188
-
189
- if follow_up:
190
- st.markdown(f"<div class='assistant-msg'><strong>HAL:</strong> {follow_up}</div>", unsafe_allow_html=True)
191
 
192
- if image_url:
193
- st.image(image_url, caption="NASA Image of the Day")
194
 
195
- st.session_state.response_ready = True
 
196
 
197
- if st.session_state.response_ready and st.session_state.follow_up:
198
- st.markdown(f"<div class='assistant-msg'><strong>HAL:</strong> {st.session_state.follow_up}</div>", unsafe_allow_html=True)
199
- st.session_state.response_ready = False
 
1
  import os
2
  import re
3
  import requests
 
4
  import streamlit as st
5
  from langchain_huggingface import HuggingFaceEndpoint
6
  from langchain_core.prompts import PromptTemplate
7
  from langchain_core.output_parsers import StrOutputParser
8
  from transformers import pipeline
9
+ from langdetect import detect
 
 
 
 
10
 
11
  # βœ… Environment Variables
12
  HF_TOKEN = os.getenv("HF_TOKEN")
13
+ NASA_API_KEY = os.getenv("NASA_API_KEY")
14
+
15
+ if not HF_TOKEN:
16
  raise ValueError("HF_TOKEN is not set. Please add it to your environment variables.")
17
 
18
+ if not NASA_API_KEY:
 
19
  raise ValueError("NASA_API_KEY is not set. Please add it to your environment variables.")
20
 
21
  # βœ… Set Up Streamlit
22
  st.set_page_config(page_title="HAL - NASA ChatBot", page_icon="πŸš€")
23
 
24
+ # βœ… Ensure Session State for Chat History
25
  if "chat_history" not in st.session_state:
26
  st.session_state.chat_history = [{"role": "assistant", "content": "Hello! How can I assist you today?"}]
27
 
28
+ # βœ… Define AI Model
29
  def get_llm_hf_inference(model_id="mistralai/Mistral-7B-Instruct-v0.3", max_new_tokens=512, temperature=0.7):
30
  return HuggingFaceEndpoint(
31
  repo_id=model_id,
32
  max_new_tokens=max_new_tokens,
33
  temperature=temperature,
34
  token=HF_TOKEN,
35
+ task="text-generation"
 
36
  )
37
 
38
+ # βœ… Generate Follow-Up Question (Preserving Format)
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
39
  def generate_follow_up(user_text):
 
 
40
  prompt_text = (
41
  f"Given the user's question: '{user_text}', generate a SHORT follow-up question in the format: "
42
+ "'Would you like to learn more about [related topic] or explore something else?'."
43
+ "Ensure it is concise and strictly follows this format."
44
  )
45
+
46
+ hf = get_llm_hf_inference(max_new_tokens=30, temperature=0.6)
47
  output = hf.invoke(input=prompt_text).strip()
48
 
 
49
  cleaned_output = re.sub(r"```|''|\"", "", output).strip()
50
 
51
+ return cleaned_output if "Would you like to learn more about" in cleaned_output else "Would you like to explore another related topic or ask about something else?"
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
52
 
53
+ # βœ… Get AI Response and Maintain Chat History
54
+ def get_response(user_text):
55
+ """Generates a response and updates chat history."""
56
+
57
+ hf = get_llm_hf_inference(max_new_tokens=512, temperature=0.9)
58
 
59
+ # Format chat history for context
60
+ filtered_history = "\n".join(f"{msg['role']}: {msg['content']}" for msg in st.session_state.chat_history)
61
 
62
+ # Create prompt
63
  prompt = PromptTemplate.from_template(
64
+ "[INST] You are a helpful AI assistant.\n\nCurrent Conversation:\n{chat_history}\n\n"
65
  "User: {user_text}.\n [/INST]\n"
66
+ "AI: Provide a detailed but concise explanation with depth. "
67
+ "Ensure a friendly, engaging tone."
68
  "\nHAL:"
69
  )
70
 
71
  chat = prompt | hf.bind(skip_prompt=True) | StrOutputParser(output_key='content')
72
+ response = chat.invoke(input=dict(user_text=user_text, chat_history=filtered_history))
73
  response = response.split("HAL:")[-1].strip() if "HAL:" in response else response.strip()
74
 
75
+ # Generate follow-up question
 
 
 
 
76
  follow_up = generate_follow_up(user_text)
77
 
78
+ # βœ… Preserve conversation history
79
+ st.session_state.chat_history.append({'role': 'user', 'content': user_text})
80
+ st.session_state.chat_history.append({'role': 'assistant', 'content': response})
81
+ st.session_state.chat_history.append({'role': 'assistant', 'content': follow_up})
 
82
 
83
+ return response, follow_up
84
 
85
+ # βœ… Chat UI
86
  st.title("πŸš€ HAL - NASA AI Assistant")
87
 
88
+ # βœ… Display Conversation History BEFORE User Input
89
+ for message in st.session_state.chat_history:
90
+ if message["role"] == "user":
91
+ st.markdown(f"<div class='user-msg'><strong>You:</strong> {message['content']}</div>", unsafe_allow_html=True)
92
+ else:
93
+ st.markdown(f"<div class='assistant-msg'><strong>HAL:</strong> {message['content']}</div>", unsafe_allow_html=True)
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
94
 
95
+ # βœ… User Input
96
  user_input = st.chat_input("Type your message here...")
97
 
98
  if user_input:
99
+ response, follow_up = get_response(user_input)
 
 
 
 
 
 
 
 
 
 
100
 
101
+ # βœ… Display AI response
102
+ st.markdown(f"<div class='assistant-msg'><strong>HAL:</strong> {response}</div>", unsafe_allow_html=True)
103
 
104
+ # βœ… Display Follow-Up Question
105
+ st.markdown(f"<div class='assistant-msg'><strong>HAL:</strong> {follow_up}</div>", unsafe_allow_html=True)
106