Update app.py

app.py

@@ -2,7 +2,8 @@ import json
 import os
 import pandas as pd
 import requests
-import threading
+from multiprocessing import Pool
+from functools import partial
 import streamlit as st
 from datasets import load_dataset, load_metric
 
@@ -32,11 +33,7 @@ def read_markdown(path):
         output = f.read()
     st.markdown(output, unsafe_allow_html=True)
 
-
-def generate_code(
-    generations, model_name, gen_prompt, max_new_tokens, temperature, seed
-):
-    # call space using its API endpoint
+def generate_code(model_name, gen_prompt, max_new_tokens, temperature, seed):
     url = (
         f"https://hf.space/embed/loubnabnl/{model_name.lower()}-subspace/+/api/predict/"
     )
@@ -44,8 +41,7 @@ def generate_code(
         url=url, json={"data": [gen_prompt, max_new_tokens, temperature, seed]}
     )
     generated_text = r.json()["data"][0]
-    generations.append({model_name: generated_text})
-
+    return generated_text
 
 def generate_code_threads(
     generations, models, gen_prompt, max_new_tokens, temperature, seed
@@ -185,24 +181,23 @@ gen_prompt = st.text_area(
     value=example_text,
     height=200,
 ).strip()
-if st.button("Generate code!", key=4):
+if st.button("Generate code!"):
     with st.spinner("Generating code..."):
-        # use threading
-        generations = []
-        generate_code_threads(
-            generations,
-            selected_models,
+        # Create a multiprocessing Pool
+        pool = Pool()
+        generate_parallel = partial(
+            generate_code,
+
             gen_prompt=gen_prompt,
             max_new_tokens=max_new_tokens,
             temperature=temperature,
             seed=seed,
         )
-        for i in range(len(generations)):
+        output = pool.map(generate_parallel, selected_models)
+        for i in range(len(output)):
             st.markdown(f"**{selected_models[i]}**")
-            for j in range(len(generations)):
-                if selected_models[i] in generations[j].keys():
-                    st.code(generations[j][selected_models[i]])
-        if len(generations) < len(selected_models):
+            st.code(output[i])
+        if len(output) < len(selected_models):
             st.markdown("<span style='color:red'>Warning: Some models run into timeout, you can try generating code using the original subspaces: [InCoder](https://huggingface.co/spaces/loubnabnl/incoder-subspace), [CodeGen](https://huggingface.co/spaces/loubnabnl/codegen-subspace), [CodeParrot](https://huggingface.co/spaces/loubnabnl/codeparrot-subspace)</span>", unsafe_allow_html=True)
 
 # Resources