parse josn output from prompt

app.py

@@ -1,6 +1,7 @@
 import spaces
 import gradio as gr
 import os
+import orjson
 import torch
 from transformers import AutoModelForSpeechSeq2Seq, AutoProcessor, pipeline, AutoModelForCausalLM, AutoTokenizer
 
@@ -36,6 +37,7 @@ def transcribe_audio(audio, model_id):
     result = pipe(audio)
     return result["text"]
 
+@spaces.GPU(duration=60)
 def proofread(text):
     if text is None:
         return "Please provide the transcribed text for proofreading."
@@ -48,7 +50,15 @@ def proofread(text):
         {"role": "user", "content": text},
     ]
     pipe = pipeline("text-generation", model="hfl/llama-3-chinese-8b-instruct-v3")
-    proofread_text = pipe(messages)
+    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
     return proofread_text
 
 with gr.Blocks() as demo:

requirements.txt

@@ -4,4 +4,5 @@ torch
 torchvision
 torchaudio
 accelerate
-datasets
+datasets
+orjson