add multiprocessing in inference and clean code

app.py

@@ -1,29 +1,27 @@
-import streamlit as st
-from transformers import AutoTokenizer, AutoModelForCausalLM
-from transformers import pipeline
-import torch
 import json
 import pandas as pd
 import requests
+from multiprocessing import Pool
+from functools import partial
+import streamlit as st
+
 
 GITHUB_CODE = "https://huggingface.co/datasets/lvwerra/github-code"
 INCODER_IMG = "https://huggingface.co/datasets/loubnabnl/repo-images/raw/main/incoder.png"
 
-@st.cache(allow_output_mutation=True)
-def load_tokenizer(model_ckpt):
-    return AutoTokenizer.from_pretrained(model_ckpt)
-
-@st.cache(allow_output_mutation=True)
-def load_model(model_ckpt):
-    model = AutoModelForCausalLM.from_pretrained(model_ckpt, low_cpu_mem_usage=True)
-    return model
-
 @st.cache()
 def load_examples():
     with open("utils/examples.json", "r") as f:
         examples = json.load(f)
     return examples
 
+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/'
+    r = requests.post(url=url, json={"data": [gen_prompt, max_new_tokens, temperature, seed]})
+    generated_text = r.json()['data'][0]
+    st.markdown(model_name)
+    st.code(generated_text)
+    
 st.set_page_config(page_icon=":laptop:", layout="wide")
 
 st.sidebar.header("Models")
@@ -84,9 +82,11 @@ elif selected_task == "Code generation":
     gen_prompt = st.text_area("Generate code with prompt:", value=example_text, height=220,).strip()
     if st.button("Generate code!"):
         with st.spinner("Generating code..."):
-            for model in selected_models:
-                url = f'https://hf.space/embed/loubnabnl/{model.lower()}-subspace/+/api/predict/'
-                r = requests.post(url=url, json={"data": [gen_prompt, max_new_tokens, temperature, seed]})
-                generated_text = r.json()['data'][0]
-                st.markdown(f"{model}")
-                st.code(generated_text)
+            # 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)
+            pool.map(generate_parallel, selected_models)