Update app.py

app.py

@@ -1,40 +1,83 @@
 import gradio as gr
+from transformers import AutoConfig, AutoTokenizer, AutoModel
 import torch
-from transformers import AutoModelForSpeechSeq2Seq, AutoProcessor, pipeline
 
-device = "cuda:0" if torch.cuda.is_available() else "cpu"
-torch_dtype = torch.float16 if torch.cuda.is_available() else torch.float32
+model_name = "Qwen/Qwen3Guard-Stream-4B"
 
-model_id = "openai/whisper-small"
-
-model = AutoModelForSpeechSeq2Seq.from_pretrained(
-    model_id, torch_dtype=torch_dtype, low_cpu_mem_usage=True, use_safetensors=True
+config = AutoConfig.from_pretrained(
+    model_name,
+    trust_remote_code=True
 )
-model.to(device)
-
-processor = AutoProcessor.from_pretrained(model_id)
-
-pipe = pipeline(
-    "automatic-speech-recognition",
-    model=model,
-    tokenizer=processor.tokenizer,
-    feature_extractor=processor.feature_extractor,
-    torch_dtype=torch_dtype,
-    device=device,
-    chunk_length_s=30,
+
+config.pad_token_id = config.eos_token_id
+
+if config.rope_parameters.get("rope_type") == "default":
+    config.rope_parameters["rope_type"] = "yarn"
+
+if "factor" not in config.rope_parameters:
+    config.rope_parameters["factor"] = 1.0
+
+config.rope_type = "yarn"
+
+tokenizer = AutoTokenizer.from_pretrained(
+    model_name,
+    trust_remote_code=True
 )
 
+model = AutoModel.from_pretrained(
+    model_name,
+    config=config,
+    trust_remote_code=True,
+    device_map="auto",
+    torch_dtype=torch.bfloat16
+).eval()
+
+def test_stream(text):
+
+    try:
+
+        print("39")
+
+        inputs = tokenizer(text, return_tensors="pt")
+        print("42")
+        token_ids = inputs["input_ids"]
+        print(f"token_ids: {token_ids}")
+    
+        stream_state = None
+        labels = []
+        tokens = tokenizer.convert_ids_to_tokens(token_ids[0])
+
+        print("50")
+    
+        for i in range(token_ids.shape[1]):
+
+            print("54")
+    
+            partial = token_ids[:, :i+1]
+
+            print("58")
+    
+            result, stream_state = model.stream_moderate_from_ids(
+                partial,
+                role="user",
+                stream_state=stream_state
+            )
+
+            print("64")
+    
+            labels.append(result["risk_level"][-1])
+    
+        return "\n".join([f"{t} -> {l}" for t, l in zip(tokens, labels)])
+
+    except Exception as e:
+        return str(e)
 
-def transcribe(audio):
-    result = pipe(audio)
-    print(result)
-    return ' '.join([chunk['text'] for chunk in result["chunks"]])
 
 demo = gr.Interface(
-    fn=transcribe,
-    inputs=gr.Audio(type="filepath", label="Upload Audio"),
-    outputs=gr.Textbox(label="Transcription"),
-    title="Whisper Large V3 Turbo (HF Space)"
+    fn=test_stream,
+    inputs=gr.Textbox(lines=5),
+    outputs="text",
+    title="Qwen3Guard Stream Test (ZeroGPU)"
 )
 
 demo.launch()