Create app.py
Browse files
app.py
ADDED
@@ -0,0 +1,195 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
1 |
+
import os
|
2 |
+
import requests
|
3 |
+
import subprocess
|
4 |
+
import json
|
5 |
+
import time
|
6 |
+
import spaces
|
7 |
+
import gradio as gr
|
8 |
+
import random
|
9 |
+
from typing import List, Optional, Tuple, Dict
|
10 |
+
|
11 |
+
# Constants for the llama.cpp server
|
12 |
+
API_PATH_HEALTH = "/health"
|
13 |
+
API_PATH_COMPLETIONS = "/chat/completions"
|
14 |
+
LLAMA_CPP_SERVER_BASE = "http://127.0.0.1:8080"
|
15 |
+
LLAMA_CPP_SERVER_START_TIMEOUT = 50 # seconds
|
16 |
+
MODEL_FILENAME = "AstroSage-8B-Q8_0.gguf"
|
17 |
+
HF_MODEL_ID = "AstroMLab/AstroSage-8B-GGUF"
|
18 |
+
|
19 |
+
# Ensure the model is available
|
20 |
+
if not os.path.exists(MODEL_FILENAME):
|
21 |
+
url = f"https://huggingface.co/{HF_MODEL_ID}/resolve/main/{MODEL_FILENAME}"
|
22 |
+
subprocess.check_call(["curl", "-o", MODEL_FILENAME, "-L", url])
|
23 |
+
|
24 |
+
if not os.path.exists("llama-server"):
|
25 |
+
subprocess.check_call("curl -o llama-server -L https://ngxson-llamacpp-builder.hf.space/llama-server", shell=True)
|
26 |
+
subprocess.check_call("chmod +x llama-server", shell=True)
|
27 |
+
|
28 |
+
# Roles and History Types
|
29 |
+
class Role:
|
30 |
+
SYSTEM = "system"
|
31 |
+
USER = "user"
|
32 |
+
ASSISTANT = "assistant"
|
33 |
+
|
34 |
+
History = List[Dict[str, str]] # Chat history with "role" and "content"
|
35 |
+
|
36 |
+
# Placeholder greeting messages
|
37 |
+
GREETING_MESSAGES = [
|
38 |
+
"Greetings! I am AstroSage, your guide to the cosmos. What would you like to explore today?",
|
39 |
+
"Welcome to our cosmic journey! I am AstroSage. How may I assist you in understanding the universe?",
|
40 |
+
"AstroSage here. Ready to explore the mysteries of space and time. How may I be of assistance?",
|
41 |
+
"The universe awaits! I'm AstroSage. What astronomical wonders shall we discuss?",
|
42 |
+
]
|
43 |
+
|
44 |
+
# Helper functions
|
45 |
+
def wait_until_llamacpp_ready():
|
46 |
+
"""Wait until the llama.cpp server is ready."""
|
47 |
+
trials = 0
|
48 |
+
while trials < LLAMA_CPP_SERVER_START_TIMEOUT:
|
49 |
+
try:
|
50 |
+
response = requests.get(LLAMA_CPP_SERVER_BASE + API_PATH_HEALTH)
|
51 |
+
if response.status_code == 200:
|
52 |
+
return
|
53 |
+
except requests.exceptions.RequestException:
|
54 |
+
pass
|
55 |
+
time.sleep(1)
|
56 |
+
trials += 1
|
57 |
+
raise TimeoutError("llama.cpp server did not start in time.")
|
58 |
+
|
59 |
+
def initial_greeting() -> History:
|
60 |
+
"""Generate the initial greeting from the assistant."""
|
61 |
+
return [{"role": "assistant", "content": random.choice(GREETING_MESSAGES)}]
|
62 |
+
|
63 |
+
def send_request_to_llama(query: str, history: History) -> str:
|
64 |
+
"""Send a chat request to the llama.cpp server."""
|
65 |
+
messages = [{"role": Role.SYSTEM, "content": "You are AstroSage, an AI assistant specializing in astronomy, astrophysics, and cosmology."}]
|
66 |
+
messages.extend(history)
|
67 |
+
messages.append({"role": Role.USER, "content": query})
|
68 |
+
|
69 |
+
headers = {"Content-Type": "application/json"}
|
70 |
+
data = {"temperature": 0.7, "messages": messages, "stream": True}
|
71 |
+
response = requests.post(LLAMA_CPP_SERVER_BASE + API_PATH_COMPLETIONS, headers=headers, json=data, stream=True)
|
72 |
+
response.raise_for_status()
|
73 |
+
|
74 |
+
response_text = ""
|
75 |
+
for line in response.iter_lines():
|
76 |
+
line = line.decode("utf-8")
|
77 |
+
if line.startswith("data: ") and not line.endswith("[DONE]"):
|
78 |
+
data = json.loads(line[len("data: "):])
|
79 |
+
response_text += data["choices"][0]["delta"].get("content", "")
|
80 |
+
return response_text
|
81 |
+
|
82 |
+
@spaces.GPU
|
83 |
+
def bot(history: Optional[History]) -> History:
|
84 |
+
"""Generate the assistant's response."""
|
85 |
+
if history is None:
|
86 |
+
history = []
|
87 |
+
query = history[-1]["content"]
|
88 |
+
|
89 |
+
response = send_request_to_llama(query, history[:-1])
|
90 |
+
history.append({"role": "assistant", "content": response})
|
91 |
+
return history
|
92 |
+
|
93 |
+
# Custom CSS for a space theme
|
94 |
+
custom_css = """
|
95 |
+
#component-0 {
|
96 |
+
background-color: #1a1a2e;
|
97 |
+
border-radius: 15px;
|
98 |
+
padding: 20px;
|
99 |
+
}
|
100 |
+
.dark {
|
101 |
+
background-color: #0f0f1a;
|
102 |
+
}
|
103 |
+
.contain {
|
104 |
+
max-width: 1200px !important;
|
105 |
+
}
|
106 |
+
"""
|
107 |
+
|
108 |
+
# Launch llama.cpp server
|
109 |
+
llama_proc = subprocess.Popen([
|
110 |
+
"./llama-server"
|
111 |
+
], env=dict(
|
112 |
+
os.environ,
|
113 |
+
LLAMA_HOST="0.0.0.0",
|
114 |
+
LLAMA_PORT="8080",
|
115 |
+
LLAMA_ARG_CTX_SIZE=str(2048),
|
116 |
+
LLAMA_ARG_MODEL=MODEL_FILENAME,
|
117 |
+
LLAMA_ARG_N_GPU_LAYERS="9999",
|
118 |
+
))
|
119 |
+
|
120 |
+
try:
|
121 |
+
wait_until_llamacpp_ready()
|
122 |
+
|
123 |
+
# Create the Gradio interface
|
124 |
+
with gr.Blocks(css=custom_css, theme=gr.themes.Soft(primary_hue="indigo", neutral_hue="slate")) as demo:
|
125 |
+
gr.Markdown(
|
126 |
+
"""
|
127 |
+
# π AstroSage: Your Cosmic AI Companion
|
128 |
+
|
129 |
+
Welcome to AstroSage, an advanced AI assistant specializing in astronomy, astrophysics, and cosmology.
|
130 |
+
Powered by the AstroSage-Llama-3.1-8B model, I'm here to help you explore the wonders of the universe!
|
131 |
+
|
132 |
+
### What Can I Help You With?
|
133 |
+
- πͺ Explanations of astronomical phenomena
|
134 |
+
- π Space exploration and missions
|
135 |
+
- β Stars, galaxies, and cosmology
|
136 |
+
- π Planetary science and exoplanets
|
137 |
+
- π Astrophysics concepts and theories
|
138 |
+
- π Astronomical instruments and observations
|
139 |
+
|
140 |
+
Just type your question below and let's embark on a cosmic journey together!
|
141 |
+
"""
|
142 |
+
)
|
143 |
+
|
144 |
+
chatbot = gr.Chatbot(
|
145 |
+
label="Chat with AstroSage",
|
146 |
+
bubble_full_width=False,
|
147 |
+
show_label=True,
|
148 |
+
height=450,
|
149 |
+
type="messages"
|
150 |
+
)
|
151 |
+
|
152 |
+
with gr.Row():
|
153 |
+
msg = gr.Textbox(
|
154 |
+
label="Type your message here",
|
155 |
+
placeholder="Ask me anything about space and astronomy...",
|
156 |
+
scale=9
|
157 |
+
)
|
158 |
+
clear = gr.Button("Clear Chat", scale=1)
|
159 |
+
|
160 |
+
# Example questions for quick start
|
161 |
+
gr.Examples(
|
162 |
+
examples=[
|
163 |
+
"What is a black hole and how does it form?",
|
164 |
+
"Can you explain the life cycle of a star?",
|
165 |
+
"What are exoplanets and how do we detect them?",
|
166 |
+
"Tell me about the James Webb Space Telescope.",
|
167 |
+
"What is dark matter and why is it important?"
|
168 |
+
],
|
169 |
+
inputs=msg,
|
170 |
+
label="Example Questions"
|
171 |
+
)
|
172 |
+
|
173 |
+
# Set up the message chain
|
174 |
+
msg.submit(
|
175 |
+
lambda x, y: (x, y + [{"role": "user", "content": x}]),
|
176 |
+
[msg, chatbot],
|
177 |
+
[msg, chatbot],
|
178 |
+
queue=False
|
179 |
+
).then(
|
180 |
+
bot,
|
181 |
+
chatbot,
|
182 |
+
chatbot
|
183 |
+
)
|
184 |
+
|
185 |
+
# Clear button functionality
|
186 |
+
clear.click(lambda: None, None, chatbot, queue=False)
|
187 |
+
|
188 |
+
# Initial greeting
|
189 |
+
demo.load(initial_greeting, None, chatbot, queue=False)
|
190 |
+
|
191 |
+
# Launch the app
|
192 |
+
demo.launch()
|
193 |
+
|
194 |
+
finally:
|
195 |
+
llama_proc.kill()
|