Upload app.py

app.py

@@ -1,70 +1,29 @@
-import os
-from threading import Thread
 import gradio as gr
-import torch
-from transformers import AutoModelForCausalLM, AutoTokenizer, TextIteratorStreamer
+from huggingface_hub import InferenceClient
 
-MODEL_ID = "clem/macron-style-qwen2.5-1.5B"
+client = InferenceClient("clem/macron-style-qwen2.5-1.5B")
 
-try:
-    import spaces
-    HAS_SPACES = True
-except ImportError:
-    HAS_SPACES = False
-
-tokenizer = AutoTokenizer.from_pretrained(MODEL_ID)
-model = AutoModelForCausalLM.from_pretrained(
-    MODEL_ID,
-    torch_dtype=torch.float16 if torch.cuda.is_available() else torch.float32,
-    device_map="auto",
-    use_cache=True,
-)
+SYSTEM_PROMPT = "You are Emmanuel Macron, President of the French Republic. Respond in his characteristic style: eloquent, diplomatic yet direct, reformist, and deeply European."
 
 
-def _generate(input_ids, max_new_tokens, temperature, top_p):
-    streamer = TextIteratorStreamer(
-        tokenizer, timeout=30.0, skip_prompt=True, skip_special_tokens=True
-    )
-    generate_kwargs = dict(
-        input_ids=input_ids.to(model.device),
-        streamer=streamer,
-        max_new_tokens=max_new_tokens,
-        do_sample=True,
+def respond(message: str, chat_history: list[dict], system_prompt: str, max_tokens: int, temperature: float, top_p: float):
+    messages = []
+    if system_prompt.strip():
+        messages.append({"role": "system", "content": system_prompt})
+    messages.extend(chat_history)
+    messages.append({"role": "user", "content": message})
+
+    response = ""
+    for chunk in client.chat_completion(
+        messages,
+        max_tokens=max_tokens,
         temperature=temperature,
         top_p=top_p,
-        eos_token_id=tokenizer.eos_token_id,
-    )
-    Thread(target=model.generate, kwargs=generate_kwargs, daemon=True).start()
-    output = ""
-    for chunk in streamer:
-        output += chunk
-        yield output
-
-
-if HAS_SPACES:
-    _generate = spaces.GPU(_generate)
-
-
-def respond(
-    message: str,
-    chat_history: list[dict],
-    system_prompt: str,
-    max_new_tokens: int,
-    temperature: float,
-    top_p: float,
-):
-    conversation = []
-    if system_prompt.strip():
-        conversation.append({"role": "system", "content": system_prompt})
-    conversation.extend(chat_history)
-    conversation.append({"role": "user", "content": message})
-
-    input_ids = tokenizer.apply_chat_template(
-        conversation,
-        add_generation_prompt=True,
-        return_tensors="pt",
-    )
-    yield from _generate(input_ids, max_new_tokens, temperature, top_p)
+        stream=True,
+    ):
+        token = chunk.choices[0].delta.content or ""
+        response += token
+        yield response
 
 
 demo = gr.ChatInterface(
@@ -72,11 +31,7 @@ demo = gr.ChatInterface(
     type="messages",
     chatbot=gr.Chatbot(height=500, type="messages"),
     additional_inputs=[
-        gr.Textbox(
-            value="You are Emmanuel Macron, President of the French Republic. Respond in his characteristic style: eloquent, diplomatic yet direct, reformist, and deeply European.",
-            label="System prompt",
-            lines=3,
-        ),
+        gr.Textbox(value=SYSTEM_PROMPT, label="System prompt", lines=3),
         gr.Slider(64, 1024, value=256, step=64, label="Max new tokens"),
         gr.Slider(0.1, 2.0, value=0.7, step=0.1, label="Temperature"),
         gr.Slider(0.1, 1.0, value=0.95, step=0.05, label="Top-p"),
@@ -91,5 +46,4 @@ demo = gr.ChatInterface(
     description="A Qwen2.5-1.5B fine-tuned to speak in the style of Emmanuel Macron. Trained on [clem/macron-style-conversations](https://hf.co/datasets/clem/macron-style-conversations).",
 )
 
-if __name__ == "__main__":
-    demo.launch()
+demo.launch()