Update app.py

app.py

@@ -1,32 +1,34 @@
-def test_eos_pad():
-    from datasets import load_dataset
-    import torch
-    from transformers import GPT2Tokenizer, GPT2LMHeadModel
+import torch
+import gradio as gr
+from transformers import AutoProcessor, AutoModelForSpeechSeq2Seq
 
-    raw_text_batch = 'a'
+# Load the processor and model
+processor = AutoProcessor.from_pretrained("NbAiLab/nb-whisper-large-verbatim")
+model = AutoModelForSpeechSeq2Seq.from_pretrained("NbAiLab/nb-whisper-large-verbatim")
 
-    tokenizer = GPT2Tokenizer.from_pretrained("gpt2")
-    # print(f'{tokenizer.eos_token=}')
-    # print(f'{tokenizer.eos_token_id=}')
-    # print(f'{tokenizer.pad_token=}')
-    # print(f'{tokenizer.pad_token_id=}')
+# Ensure the pad token is set
+if processor.tokenizer.pad_token_id is None:
+    processor.tokenizer.pad_token = processor.tokenizer.eos_token
 
-    # print(f'{raw_text_batch=}')
-    # tokenize_batch = tokenizer(raw_text_batch, padding="max_length", max_length=5, truncation=True, return_tensors="pt")
-    # print(f'{tokenize_batch=}')
+# Move the model to the appropriate device
+device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+model = model.to(device)
 
-    if tokenizer.pad_token_id is None:
-        tokenizer.pad_token = tokenizer.eos_token
-    probe_network = GPT2LMHeadModel.from_pretrained("gpt2")
-    device = torch.device(f"cuda:{0}" if torch.cuda.is_available() else "cpu")
-    probe_network = probe_network.to(device)
+def transcribe_text(raw_text):
+    # Tokenize the input text
+    tokenize_batch = processor.tokenizer(raw_text, padding="max_length", max_length=5, truncation=True, return_tensors="pt")
+    
+    # Output the tokenized result for inspection
+    return tokenize_batch
 
-    print(f'{tokenizer.eos_token=}')
-    print(f'{tokenizer.eos_token_id=}')
-    print(f'{tokenizer.pad_token=}')
-    print(f'{tokenizer.pad_token_id=}')
+# Gradio interface
+iface = gr.Interface(
+    fn=transcribe_text,
+    inputs=gr.inputs.Textbox(lines=2, placeholder="Enter text here..."),
+    outputs="json",
+    title="Whisper Model Tokenization",
+    description="Test the EOS and PAD tokens for NbAiLab/nb-whisper-large-verbatim model."
+)
 
-    print(f'{raw_text_batch=}')
-    tokenize_batch = tokenizer(raw_text_batch, padding="max_length", max_length=5, truncation=True, return_tensors="pt")
-    print(f'{tokenize_batch=}')
-    print('Done')
+if __name__ == "__main__":
+    iface.launch()