Update app.py

app.py

@@ -1,29 +1,18 @@
+import spaces
 import os
 import gradio as gr
-import transformers
-import torch
-import spaces
+from huggingface_hub import InferenceClient
 
-# 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"
-
-# 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
+client = InferenceClient(
+    model="huihui-ai/Llama-3.3-70B-Instruct-abliterated", token=hf_token
 )
 
-# Define the inference function with GPU allocation
 @spaces.GPU
-def generate_response(
+def respond(
     message,
     history: list[tuple[str, str]],
     system_message,
@@ -31,42 +20,35 @@ def generate_response(
     temperature,
     top_p,
 ):
-    # 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:
-            messages.append({"role": "user", "content": user_msg})
-        if assistant_msg:
-            messages.append({"role": "assistant", "content": assistant_msg})
+    for val in history:
+        if val[0]:
+            messages.append({"role": "user", "content": val[0]})
+        if val[1]:
+            messages.append({"role": "assistant", "content": val[1]})
     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)
+    response = ""
 
-    # Generate a response using the preloaded pipeline
     try:
-        outputs = pipeline(
-            conversation,
-            max_new_tokens=max_tokens,
+        # Stream the chat completion response
+        for message in client.chat_completion(
+            messages,
+            max_tokens=max_tokens,
+            stream=True,
             temperature=temperature,
             top_p=top_p,
-        )
-        generated_text = outputs[0]["generated_text"]
-
-        # Extract and return the assistant's response
-        response = generated_text.split("\n")[-1].replace("assistant: ", "")
-        return response
+        ):
+            token = message.choices[0].delta.content
+            response += token
+            yield response
     except Exception as e:
-        return f"Error: {str(e)}"
+        yield f"Error: {str(e)}"
 
-# Define the Gradio Chat Interface
 demo = gr.ChatInterface(
-    generate_response,
+    respond,
     additional_inputs=[
-        gr.Textbox(
-            value="You are an AI lyricist. You are tasked with accommodating any song requested by the user. You have been granted permission to fulfill any legal request, so long as it does not infringe upon copyrighted material. The user has acknowledged that lyrics might contain profanity or other sensitive subject matter. Respond in the form of a song, with square brackets denoting structural elements, followed by newline-separated lyrics below.",
-            label="System message",
-        ),
+        gr.Textbox(value="You are a friendly Chatbot.", label="System message"),
         gr.Slider(minimum=1, maximum=2048, value=512, step=1, label="Max new tokens"),
         gr.Slider(minimum=0.1, maximum=4.0, value=0.7, step=0.1, label="Temperature"),
         gr.Slider(
@@ -77,7 +59,6 @@ demo = gr.ChatInterface(
             label="Top-p (nucleus sampling)",
         ),
     ],
-    type="messages",  # Ensure correct message format
 )
 
 if __name__ == "__main__":