Create app.py

app.py

@@ -0,0 +1,84 @@
+import gradio as gr
+from transformers import AutoModelForCausalLM, AutoTokenizer, TextIteratorStreamer
+import torch
+from threading import Thread
+import time
+
+model_name = "deepseek-ai/DeepSeek-R1-Distill-Qwen-1.5B"
+tokenizer = AutoTokenizer.from_pretrained(model_name)
+tokenizer.chat_template = "{% for message in messages %}<|im_start|>{{ message.role }}\n{{ message.content }}<|im_end|>\n{% endfor %}<|im_start|>assistant\n"
+
+model = AutoModelForCausalLM.from_pretrained(
+    model_name,
+    torch_dtype=torch.float32,
+    device_map="cpu",
+    low_cpu_mem_usage=True
+).to('cpu')
+class deepstreamer(TextIteratorStreamer):
+    def __init__(self, tokenizer):
+        super().__init__(tokenizer, skip_prompt=True, skip_special_tokens=True)
+        self.token_count = 0
+        self.start_time = None
+
+    def put(self, value):
+        if self.start_time is None:
+            self.start_time = time.time()
+        self.token_count += 1
+        return super().put(value)
+
+    def get_tps(self):
+        if self.start_time is None:
+            return 0
+        return self.token_count / (time.time() - self.start_time)
+
+def format_response(text, tps=None): #token per second
+    return f"{text}\n\n**Tokens per second:** {tps:.2f}" if tps else text
+def chat_response(message, history, max_tokens=512):
+    messages = []
+    for human, assistant in history:
+        messages.append({"role": "user", "content": human})
+        messages.append({"role": "assistant", "content": assistant})
+    messages.append({"role": "user", "content": message})
+    formatted_input = tokenizer.apply_chat_template(
+        messages,
+        tokenize=False,
+        add_generation_prompt=True
+    )    
+    inputs = tokenizer(formatted_input, return_tensors="pt").to('cpu')
+    streamer = deepstreamer(tokenizer)
+    generation_kwargs = dict(
+        inputs,
+        streamer=streamer,
+        max_new_tokens=max_tokens,
+        do_sample=True,
+        temperature=0.7,
+        top_p=0.9,
+        eos_token_id=tokenizer.eos_token_id
+    )
+    thread = Thread(target=model.generate, kwargs=generation_kwargs)
+    thread.start()
+    partial_response = ""
+    try:
+        for token in streamer:
+            partial_response += token
+            yield format_response(partial_response, streamer.get_tps())
+        
+        final_tps = streamer.token_count / (time.time() - streamer.start_time)
+        yield format_response(partial_response, final_tps)
+    finally:
+        thread.join()
+        
+demo = gr.ChatInterface(
+    fn=chat_response,
+    title="DeepSeek-R1-Distill-Qwen-1.5B on CPU",
+    description="Runnig on CPU so expect less tokens",
+    examples=[
+        "Discuss the future of renewable energy",
+        "What's the history of the Roman Empire?",
+        "What's the capital of China?",
+        "Tell me a fun fact about space"
+    ]
+)
+
+if __name__ == "__main__":
+    demo.queue().launch()