Merge branch 'main' of https://huggingface.co/spaces/Ruurd/radiolm

app.py

@@ -1,6 +1,8 @@
 import os
 import torch
 import time
+import torch
+import time
 import gradio as gr
 import spaces
 from transformers import AutoTokenizer, AutoModelForCausalLM, TextIteratorStreamer
@@ -12,6 +14,51 @@ import threading
 from transformers import TextIteratorStreamer
 import queue
 
+class RichTextStreamer(TextIteratorStreamer):
+    def __init__(self, tokenizer, prompt_len=0, **kwargs):
+        super().__init__(tokenizer, **kwargs)
+        self.token_queue = queue.Queue()
+        self.prompt_len = prompt_len
+        self.count = 0
+
+    def put(self, value):
+        if isinstance(value, torch.Tensor):
+            token_ids = value.view(-1).tolist()
+        elif isinstance(value, list):
+            token_ids = value
+        else:
+            token_ids = [value]
+
+        for token_id in token_ids:
+            self.count += 1
+            if self.count <= self.prompt_len:
+                continue  # skip prompt tokens
+            token_str = self.tokenizer.decode([token_id], **self.decode_kwargs)
+            is_special = token_id in self.tokenizer.all_special_ids
+            self.token_queue.put({
+                "token_id": token_id,
+                "token": token_str,
+                "is_special": is_special
+            })
+
+    def __iter__(self):
+        while True:
+            try:
+                token_info = self.token_queue.get(timeout=self.timeout)
+                yield token_info
+            except queue.Empty:
+                if self.end_of_generation.is_set():
+                    break
+
+from transformers import AutoTokenizer, AutoModelForCausalLM, TextIteratorStreamer
+import threading
+
+from transformers import TextIteratorStreamer
+import threading
+
+from transformers import TextIteratorStreamer
+import queue
+
 class RichTextStreamer(TextIteratorStreamer):
     def __init__(self, tokenizer, prompt_len=0, **kwargs):
         super().__init__(tokenizer, **kwargs)
@@ -194,10 +241,13 @@ def add_user_message(user_input, history):
 
 # Curated models
 model_choices = [
+    "meta-llama/Llama-3.2-3B-Instruct",
     "meta-llama/Llama-3.2-3B-Instruct",
     "deepseek-ai/DeepSeek-R1-Distill-Llama-8B",
     "google/gemma-7b",
     "mistralai/Mistral-Small-3.1-24B-Instruct-2503"
+    "google/gemma-7b",
+    "mistralai/Mistral-Small-3.1-24B-Instruct-2503"
 ]
 
 with gr.Blocks() as demo: