Update app.py

app.py

@@ -2,8 +2,7 @@ import json
 import os
 import pandas as pd
 import requests
-from multiprocessing import Pool
-from functools import partial
+import threading
 import streamlit as st
 from datasets import load_dataset, load_metric
 
@@ -33,7 +32,11 @@ def read_markdown(path):
         output = f.read()
     st.markdown(output, unsafe_allow_html=True)
 
-def generate_code(model_name, gen_prompt, max_new_tokens, temperature, seed):
+
+def generate_code(
+    generations, model_name, gen_prompt, max_new_tokens, temperature, seed
+):
+    # call space using its API endpoint
     url = (
         f"https://hf.space/embed/loubnabnl/{model_name.lower()}-subspace/+/api/predict/"
     )
@@ -41,7 +44,8 @@ def generate_code(model_name, gen_prompt, max_new_tokens, temperature, seed):
         url=url, json={"data": [gen_prompt, max_new_tokens, temperature, seed]}
     )
     generated_text = r.json()["data"][0]
-    return generated_text
+    generations.append(generated_text)
+
 
 def generate_code_threads(
     generations, models, gen_prompt, max_new_tokens, temperature, seed
@@ -181,23 +185,22 @@ gen_prompt = st.text_area(
     value=example_text,
     height=200,
 ).strip()
-if st.button("Generate code!"):
+if st.button("Generate code!", key=4):
     with st.spinner("Generating code..."):
-        # Create a multiprocessing Pool
-        pool = Pool()
-        generate_parallel = partial(
-            generate_code,
-
+        # use threading
+        generations = []
+        generate_code_threads(
+            generations,
+            selected_models,
             gen_prompt=gen_prompt,
             max_new_tokens=max_new_tokens,
             temperature=temperature,
             seed=seed,
         )
-        output = pool.map(generate_parallel, selected_models)
-        for i in range(len(output)):
+        for i in range(len(generations)):
             st.markdown(f"**{selected_models[i]}**")
-            st.code(output[i])
-        if len(output) < len(selected_models):
+            st.code(generations[i])
+        if len(generations) < 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