import os import re import requests import torch import streamlit as st from langchain_huggingface import HuggingFaceEndpoint from langchain_core.prompts import PromptTemplate from langchain_core.output_parsers import StrOutputParser from transformers import pipeline from langdetect import detect # Ensure this package is installed # ✅ Check for GPU or Default to CPU device = "cuda" if torch.cuda.is_available() else "cpu" print(f"✅ Using device: {device}") # Debugging info # ✅ Environment Variables HF_TOKEN = os.getenv("HF_TOKEN") if HF_TOKEN is None: raise ValueError("HF_TOKEN is not set. Please add it to your environment variables.") NASA_API_KEY = os.getenv("NASA_API_KEY") if NASA_API_KEY is None: raise ValueError("NASA_API_KEY is not set. Please add it to your environment variables.") # ✅ Set Up Streamlit st.set_page_config(page_title="HAL - NASA ChatBot", page_icon="🚀") # ✅ Initialize Session State Variables (Ensuring Chat History Persists) if "chat_history" not in st.session_state: st.session_state.chat_history = [{"role": "assistant", "content": "Hello! How can I assist you today?"}] if "response_ready" not in st.session_state: st.session_state.response_ready = False if "follow_up" not in st.session_state: st.session_state.follow_up = "" # ✅ Initialize Hugging Face Model (Explicitly Set to CPU/GPU) def get_llm_hf_inference(model_id="meta-llama/Llama-2-7b-chat-hf", max_new_tokens=800, temperature=0.8): #mistralai/Mistral-7B-Instruct-v0.3 return HuggingFaceEndpoint( repo_id=model_id, max_new_tokens=max_new_tokens, temperature=temperature, token=HF_TOKEN, task="text-generation", device=-1 if device == "cpu" else 0 # ✅ Force CPU (-1) or GPU (0) ) # ✅ NASA API Function def get_nasa_apod(): url = f"https://api.nasa.gov/planetary/apod?api_key={NASA_API_KEY}" response = requests.get(url) if response.status_code == 200: data = response.json() return data.get("url", ""), data.get("title", ""), data.get("explanation", "") return "", "NASA Data Unavailable", "I couldn't fetch data from NASA right now." # ✅ Sentiment Analysis (Now Uses Explicit Device) sentiment_analyzer = pipeline( "sentiment-analysis", model="distilbert/distilbert-base-uncased-finetuned-sst-2-english", device=-1 if device == "cpu" else 0 # ✅ Force CPU (-1) or GPU (0) ) def analyze_sentiment(user_text): result = sentiment_analyzer(user_text)[0] return result['label'] # ✅ Intent Detection def predict_action(user_text): if "NASA" in user_text.lower() or "space" in user_text.lower(): return "nasa_info" return "general_query" # ✅ Ensure English Responses def ensure_english(text): try: detected_lang = detect(text) if detected_lang != "en": return "⚠️ Sorry, I only respond in English. Can you rephrase your question?" except: return "⚠️ Language detection failed. Please ask your question again." return text # ✅ Follow-Up Question Generation def generate_follow_up(user_text): """Generates a structured follow-up question in a concise format.""" prompt_text = ( f"Given the user's question: '{user_text}', generate a SHORT follow-up question in the format: " "'Would you like to learn more about [related topic] or explore something else?'. " "Ensure it's concise and structured exactly as requested without extra commentary." ) hf = get_llm_hf_inference(max_new_tokens=30, temperature=0.8) # 🔥 Lower temp for consistency output = hf.invoke(input=prompt_text).strip() # ✅ Extract the relevant part using regex to remove unwanted symbols or truncations cleaned_output = re.sub(r"```|''|\"", "", output).strip() # ✅ Ensure output is formatted correctly if "Would you like to learn more about" not in cleaned_output: cleaned_output = "Would you like to explore another related topic or ask about something else?" return cleaned_output # ✅ Main Response Function def get_response(system_message, chat_history, user_text, max_new_tokens=800): action = predict_action(user_text) # ✅ Handle NASA-Specific Queries if action == "nasa_info": nasa_url, nasa_title, nasa_explanation = get_nasa_apod() response = f"**{nasa_title}**\n\n{nasa_explanation}" follow_up = generate_follow_up(user_text) chat_history.extend([ {'role': 'user', 'content': user_text}, {'role': 'assistant', 'content': response}, {'role': 'assistant', 'content': follow_up} ]) return response, follow_up, chat_history, nasa_url # ✅ Invoke Hugging Face Model hf = get_llm_hf_inference(max_new_tokens=max_new_tokens, temperature=0.9) filtered_history = "\n".join(f"{msg['role']}: {msg['content']}" for msg in chat_history) prompt = PromptTemplate.from_template( "[INST] You are a helpful AI assistant.\n\nCurrent Conversation:\n{chat_history}\n\n" "User: {user_text}.\n [/INST]\n" "AI: Provide a detailed explanation with depth. Use a conversational tone. " "🚨 Answer **only in English**." "Ensure a friendly, engaging tone." "\nHAL:" ) chat = prompt | hf.bind(skip_prompt=True) | StrOutputParser(output_key='content') response = chat.invoke(input=dict(system_message=system_message, user_text=user_text, chat_history=filtered_history)) response = response.split("HAL:")[-1].strip() if "HAL:" in response else response.strip() response = ensure_english(response) if not response: response = "I'm sorry, but I couldn't generate a response. Can you rephrase your question?" follow_up = generate_follow_up(user_text) # ✅ Preserve conversation history st.session_state.chat_history.append({'role': 'user', 'content': user_text}) st.session_state.chat_history.append({'role': 'assistant', 'content': response}) st.session_state.chat_history.append({'role': 'assistant', 'content': follow_up}) return response, follow_up, chat_history, None # ✅ Streamlit UI st.title("🚀 HAL - NASA AI Assistant") # ✅ Justify all chatbot responses st.markdown(""" """, unsafe_allow_html=True) # ✅ Reset Chat Button if st.sidebar.button("Reset Chat"): st.session_state.chat_history = [{"role": "assistant", "content": "Hello! How can I assist you today?"}] st.session_state.response_ready = False st.session_state.follow_up = "" # ✅ Chat UI user_input = st.chat_input("Type your message here...") if user_input: response, follow_up, st.session_state.chat_history, image_url = get_response( system_message="You are a helpful AI assistant.", user_text=user_input, chat_history=st.session_state.chat_history ) if response: st.markdown(f"
HAL: {response}
", unsafe_allow_html=True) if follow_up: st.markdown(f"
HAL: {follow_up}
", unsafe_allow_html=True) if image_url: st.image(image_url, caption="NASA Image of the Day") st.session_state.response_ready = True if st.session_state.response_ready and st.session_state.follow_up: st.markdown(f"
HAL: {st.session_state.follow_up}
", unsafe_allow_html=True) st.session_state.response_ready = False