try to fix runtime error on HF

app.py

@@ -16,7 +16,7 @@ def transcribe_audio(audio):
     if transcribe_model is None:
         return "Please load the transcription model first."
 
-    device = "cuda:0" if torch.cuda.is_available() else "mps" if torch.backends.mps.is_available() else "cpu"
+    device = "cuda:0" if torch.cuda.is_available() else "cpu"
     torch_dtype = torch.float16 if torch.cuda.is_available() else torch.float32
 
     processor = AutoProcessor.from_pretrained(transcribe_model)
@@ -44,28 +44,24 @@ def proofread(text):
     if proofread_model is None:
         return "Please load the proofreading model first."
     
-    device = "cuda:0" if torch.cuda.is_available() else "mps" if torch.backends.mps.is_available() else "cpu"
+    device = "cuda:0" if torch.cuda.is_available() else "cpu"
     torch_dtype = torch.float16 if torch.cuda.is_available() else torch.float32
     
     messages = [
         {"role": "system", "content": "用繁體中文語體文整理這段文字，在最後加上整段文字的重點。"},
         {"role": "user", "content": text},
     ]
-    pipe = pipeline("text-generation", model=proofread_model)
-    llm_output = pipe(messages)
     
-    # Extract the generated text
-    generated_text = llm_output[0]['generated_text']
-
-    # Extract the assistant's content
-    assistant_content = next(item['content'] for item in generated_text if item['role'] == 'assistant')
-
-    proofread_text = assistant_content
+    inputs = proofread_model.tokenizer(messages, return_tensors="tf", padding=True)
+    outputs = proofread_model.generate(**inputs)
+    proofread_text = proofread_model.tokenizer.decode(outputs[0], skip_special_tokens=True)
+    
     return proofread_text
 
+@spaces.GPU(duration=120)
 def load_models(transcribe_model_id, proofread_model_id):
     global transcribe_model, proofread_model
-    device = "cuda:0" if torch.cuda.is_available() else "mps" if torch.backends.mps.is_available() else "cpu"
+    device = "cuda:0" if torch.cuda.is_available() else "cpu"
     torch_dtype = torch.float16 if torch.cuda.is_available() else torch.float32
     
     transcribe_model = AutoModelForSpeechSeq2Seq.from_pretrained(