update app

app.py

@@ -37,20 +37,20 @@ tokenizer1 = load_tokenizer("lvwerra/codeparrot")
 model1 = load_model("lvwerra/codeparrot")
 tokenizer2 = load_tokenizer("facebook/incoder-1B")
 model2 = load_model("facebook/incoder-1B")
-#tokenizer3 = load_tokenizer("facebook/opt-1.3b")
-#model3 = load_model("facebook/opt-1.3b")
+tokenizer3 = load_tokenizer("facebook/opt-1.3b")
+model3 = load_model("facebook/opt-1.3b")
 pipelines = {}
-for model in models:
-    if model == "CodeParrot":
-        pipelines[model] = pipeline("text-generation", model=model1, tokenizer=tokenizer1)
-    elif model == "InCoder":
+for element in models:
+    if element == "CodeParrot":
+        pipelines[element] = pipeline("text-generation", model=model1, tokenizer=tokenizer1)
+    elif element == "InCoder":
         tokenizer = load_tokenizer("facebook/incoder-1B")
         model = load_model("facebook/incoder-1B")
-        pipelines[model] = pipeline("text-generation", model=model2, tokenizer=tokenizer2)
-    #else:
-    #    tokenizer = load_tokenizer("facebook/opt-1.3b")
-    #    model = load_model("facebook/opt-1.3b")
-    #    pipelines[model] = pipeline("text-generation", model=model3, tokenizer=tokenizer3)
+        pipelines[element] = pipeline("text-generation", model=model2, tokenizer=tokenizer2)
+    else:
+        tokenizer = load_tokenizer("facebook/opt-1.3b")
+        model = load_model("facebook/opt-1.3b")
+        pipelines[element] = pipeline("text-generation", model=model3, tokenizer=tokenizer3)
         
 examples = load_examples()
 example_names = [example["name"] for example in examples]
@@ -94,9 +94,7 @@ elif selected_task == "Code generation":
     if st.button("Generate code!"):
         with st.spinner("Generating code..."):
             for model in selected_models:
-                if model != "OPT":
-                    st.markdown(f"{len(pipelines)} model is {model} keys {pipelines.keys()}:")
-                    pipe = pipelines[model]
-                    generated_text = pipe(gen_prompt, **gen_kwargs)[0]['generated_text']
-                    st.markdown(f"{model}:")
-                    st.code(generated_text)
+                pipe = pipelines[model]
+                generated_text = pipe(gen_prompt, **gen_kwargs)[0]['generated_text']
+                st.markdown(f"{model}:")
+                st.code(generated_text)