Update Demo.py

Demo.py

@@ -65,21 +65,13 @@ def create_pipeline(model):
 def fit_data(pipeline, fed_data):
     """Fit the data into the pipeline and return the transcription."""
     data, sampling_rate = librosa.load(fed_data, sr=16000)
-    data = [float(x) for x in data]
-    
-    schema = StructType([
-        StructField("audio_content", ArrayType(FloatType())),
-        StructField("sampling_rate", LongType())
-    ])
-
-    df = pd.DataFrame({
-        "audio_content": [data],
-        "sampling_rate": [sampling_rate]
-    })
+    data = data.tolist()
+    spark_df = spark.createDataFrame([[data]], ["audio_content"])
 
-    spark_df = spark.createDataFrame(df, schema)
-    pipeline_df = pipeline.fit(spark_df).transform(spark_df)
-    return pipeline_df.select("text.result")
+    model = pipeline.fit(spark_df)
+    lp = LightPipeline(model)
+    lp_result = lp.fullAnnotate(data)[0]
+    return lp_result
 
 def save_uploadedfile(uploadedfile, path):
     """Save the uploaded file to the specified path."""
@@ -119,7 +111,7 @@ st.sidebar.markdown("""
 """, unsafe_allow_html=True)
 
 # Load examples
-AUDIO_FILE_PATH = "inputs"
+AUDIO_FILE_PATH = "/content/Wav2Vec2ForCTC/inputs"
 audio_files = sorted(os.listdir(AUDIO_FILE_PATH))
 
 selected_audio = st.selectbox("Select an audio", audio_files)
@@ -146,4 +138,4 @@ pipeline = create_pipeline(model)
 output = fit_data(pipeline, selected_audio)
 
 st.subheader(f"Transcription:")
-st.markdown(f"**{output[0]}**")
+st.markdown(f"{(output['text'][0].result).title()}")