Update app.py

app.py

@@ -2,6 +2,7 @@ from threading import Thread
 import gradio as gr
 import random
 import torch
+import spaces
 from transformers import (
     AutoModelForCausalLM,
     AutoTokenizer,
@@ -10,8 +11,8 @@ from transformers import (
 )
 
 # Constants for the model and configuration
-MODEL_ID = "universeTBD/astrollama"
-WINDOW_SIZE = 4096
+MODEL_ID = "AstroMLab/AstroSage-8B"
+WINDOW_SIZE = 2048
 DEVICE = "cuda"
 
 # Load model configuration, tokenizer, and model
@@ -35,7 +36,6 @@ GREETING_MESSAGES = [
     "The universe awaits! I'm AstroSage. What astronomical wonders shall we discuss?",
 ]
 
-
 def generate_text(prompt: str, history: list, max_new_tokens=512, temperature=0.7, top_p=0.95, top_k=50):
     """
     Generate a response using the transformer model.
@@ -86,7 +86,7 @@ def user(user_message, history):
         history = []
     return "", history + [{"role": "user", "content": user_message}]
 
-
+@spaces.GPU
 def bot(history):
     """
     Generate the bot's response based on the history.