CCockrum commited on
Commit
789726d
Β·
verified Β·
1 Parent(s): 53b4dae

Create app.py

Browse files
Files changed (1) hide show
  1. app.py +128 -0
app.py ADDED
@@ -0,0 +1,128 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
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="meta-llama/Llama-2-7b-chat-hf", max_new_tokens=800, temperature=0.3):
34
+ return HuggingFaceEndpoint(
35
+ repo_id=model_id,
36
+ max_new_tokens=max_new_tokens,
37
+ temperature=temperature, # πŸ”₯ Lowered temperature for more factual and structured responses
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
+ # βœ… Ensure English Responses
44
+ def ensure_english(text):
45
+ try:
46
+ detected_lang = detect(text)
47
+ if detected_lang != "en":
48
+ return "⚠️ Sorry, I only respond in English. Can you rephrase your question?"
49
+ except:
50
+ return "⚠️ Language detection failed. Please ask your question again."
51
+ return text
52
+
53
+ # βœ… Main Response Function (Fixing Repetition & Context)
54
+ def get_response(system_message, chat_history, user_text, max_new_tokens=800):
55
+ # βœ… Ensure conversation history is included correctly
56
+ filtered_history = "\n".join(
57
+ f"{msg['role'].capitalize()}: {msg['content']}"
58
+ for msg in chat_history[-5:] # βœ… Only keep the last 5 exchanges to prevent overflow
59
+ )
60
+
61
+ prompt = PromptTemplate.from_template(
62
+ "[INST] You are a highly knowledgeable AI assistant. Answer concisely, avoid repetition, and structure responses well."
63
+ "\n\nCONTEXT:\n{chat_history}\n"
64
+ "\nLATEST USER INPUT:\nUser: {user_text}\n"
65
+ "\n[END CONTEXT]\n"
66
+ "Assistant:"
67
+ )
68
+
69
+ # βœ… Invoke Hugging Face Model
70
+ hf = get_llm_hf_inference(max_new_tokens=max_new_tokens, temperature=0.3) # πŸ”₯ Lowered temperature
71
+ chat = prompt | hf.bind(skip_prompt=True) | StrOutputParser(output_key='content')
72
+
73
+ response = chat.invoke(input=dict(system_message=system_message, user_text=user_text, chat_history=filtered_history))
74
+
75
+ # Clean up the response - remove any "HAL:" prefix if present
76
+ response = response.split("HAL:")[-1].strip() if "HAL:" in response else response.strip()
77
+ response = ensure_english(response)
78
+
79
+ if not response:
80
+ response = "I'm sorry, but I couldn't generate a response. Can you rephrase your question?"
81
+
82
+ # βœ… Update conversation history
83
+ chat_history.append({'role': 'user', 'content': user_text})
84
+ chat_history.append({'role': 'assistant', 'content': response})
85
+
86
+ # βœ… Keep only last 10 exchanges to prevent unnecessary repetition
87
+ return response, chat_history[-10:]
88
+
89
+ # βœ… Streamlit UI
90
+ st.title("πŸš€ HAL - NASA AI Assistant")
91
+
92
+ # βœ… Justify all chatbot responses
93
+ st.markdown("""
94
+ <style>
95
+ .user-msg, .assistant-msg {
96
+ padding: 11px;
97
+ border-radius: 10px;
98
+ margin-bottom: 5px;
99
+ width: fit-content;
100
+ max-width: 80%;
101
+ text-align: justify;
102
+ }
103
+ .user-msg { background-color: #696969; color: white; }
104
+ .assistant-msg { background-color: #333333; color: white; }
105
+ .container { display: flex; flex-direction: column; align-items: flex-start; }
106
+ @media (max-width: 600px) { .user-msg, .assistant-msg { font-size: 16px; max-width: 100%; } }
107
+ </style>
108
+ """, unsafe_allow_html=True)
109
+
110
+ # βœ… Chat UI
111
+ user_input = st.chat_input("Type your message here...")
112
+
113
+ if user_input:
114
+ # Get response and update chat history
115
+ response, st.session_state.chat_history = get_response(
116
+ system_message="You are a helpful AI assistant.",
117
+ user_text=user_input,
118
+ chat_history=st.session_state.chat_history
119
+ )
120
+
121
+ # βœ… Display chat history (ONLY display from history, not separately)
122
+ st.markdown("<div class='container'>", unsafe_allow_html=True)
123
+ for message in st.session_state.chat_history:
124
+ if message["role"] == "user":
125
+ st.markdown(f"<div class='user-msg'><strong>You:</strong> {message['content']}</div>", unsafe_allow_html=True)
126
+ else:
127
+ st.markdown(f"<div class='assistant-msg'><strong>HAL:</strong> {message['content']}</div>", unsafe_allow_html=True)
128
+ st.markdown("</div>", unsafe_allow_html=True)