Update app.py

app.py

@@ -4,15 +4,24 @@ import transformers
 import torch
 import spaces
 
+# Load Hugging Face token from environment variables
 hf_token = os.getenv("HF_TOKEN")
 if not hf_token:
     raise ValueError("HF_TOKEN is not set in environment variables!")
 
+# Model ID
 model_id = "huihui-ai/Llama-3.3-70B-Instruct-abliterated"
 
-pipeline = None
-
+# Initialize the pipeline at startup
+pipeline = transformers.pipeline(
+    "text-generation",
+    model=model_id,
+    use_auth_token=hf_token,
+    model_kwargs={"torch_dtype": torch.bfloat16},  # Optimize memory usage
+    device_map="auto",  # Automatically map to available GPUs
+)
 
+# Define the inference function with GPU allocation
 @spaces.GPU
 def generate_response(
     message,
@@ -22,17 +31,7 @@ def generate_response(
     temperature,
     top_p,
 ):
-    global pipeline
-
-    if pipeline is None:
-        pipeline = transformers.pipeline(
-            "text-generation",
-            model=model_id,
-            use_auth_token=hf_token,
-            model_kwargs={"torch_dtype": torch.bfloat16},
-            device_map="auto",
-        )
-
+    # Combine system, history, and user messages into a formatted input string
     messages = [{"role": "system", "content": system_message}]
     for user_msg, assistant_msg in history:
         if user_msg:
@@ -40,8 +39,11 @@ def generate_response(
         if assistant_msg:
             messages.append({"role": "assistant", "content": assistant_msg})
     messages.append({"role": "user", "content": message})
+
+    # Format the conversation as a single string
     conversation = "\n".join(f"{m['role']}: {m['content']}" for m in messages)
 
+    # Generate a response using the preloaded pipeline
     try:
         outputs = pipeline(
             conversation,
@@ -51,11 +53,13 @@ def generate_response(
         )
         generated_text = outputs[0]["generated_text"]
 
+        # Extract and return the assistant's response
         response = generated_text.split("\n")[-1].replace("assistant: ", "")
         return response
     except Exception as e:
         return f"Error: {str(e)}"
 
+# Define the Gradio Chat Interface
 demo = gr.ChatInterface(
     generate_response,
     additional_inputs=[