Tijmen2 commited on
Commit
8e164af
Β·
verified Β·
1 Parent(s): ec23c0a

Create app.py

Browse files
Files changed (1) hide show
  1. app.py +195 -0
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()