Update app.py

app.py

@@ -6,24 +6,25 @@ import streamlit as st
 from PIL import Image
 
 from transformers import RobertaTokenizer, T5ForConditionalGeneration
+from transformers import pipeline
 
-@st.cache_resource
-def load_model(model_name):
-    # load model
-    model = T5ForConditionalGeneration.from_pretrained("hanbin/MaMaL-Com")
-    # load tokenizer
-    tokenizer = RobertaTokenizer.from_pretrained("hanbin/MaMaL-Com")
-    return model,tokenizer
+# @st.cache_resource
+# def load_model(model_name):
+#     # load model
+#     model = T5ForConditionalGeneration.from_pretrained("E:\DenseRetrievalGroup\卢帅学长ckpt\py150_model\checkpoint")
+#     # load tokenizer
+#     tokenizer = RobertaTokenizer.from_pretrained("E:\DenseRetrievalGroup\卢帅学长ckpt\py150_model\checkpoint")
+#     return model,tokenizer
 
 
 
 
 
-def main(model,tokenizer):
+def main():
     # `st.set_page_config` is used to display the default layout width, the title of the app, and the emoticon in the browser tab.
 
     st.set_page_config(
-        layout="centered", page_title="MaMaL-Gen Demo(代码生成)", page_icon="❄️"
+        layout="centered", page_title="MaMaL-Com Demo(代码补全)", page_icon="❄️"
     )
 
     c1, c2 = st.columns([0.32, 2])
@@ -47,23 +48,29 @@ def main(model,tokenizer):
 
     st.sidebar.image("images/panda.png",width=270)
 
-    st.sidebar.write("")
+    st.sidebar.markdown("---")
 
+    st.sidebar.write(
+    """
+    # 使用方法：
+    在【输入】文本框输入未完成的代码，点击【补全】按钮,即会显示补全的代码。
+    """
+    )
     # For elements to be displayed in the sidebar, we need to add the sidebar element in the widget.
 
     # We create a text input field for users to enter their API key.
 
-    API_KEY = st.sidebar.text_input(
-        "Enter your HuggingFace API key",
-        help="Once you created you HuggingFace account, you can get your free API token in your settings page: https://huggingface.co/settings/tokens",
-        type="password",
-    )
-
-    # Adding the HuggingFace API inference URL.
-    API_URL = "https://api-inference.huggingface.co/models/valhalla/distilbart-mnli-12-3"
-
-    # Now, let's create a Python dictionary to store the API headers.
-    headers = {"Authorization": f"Bearer {API_KEY}"}
+    # API_KEY = st.sidebar.text_input(
+    #     "Enter your HuggingFace API key",
+    #     help="Once you created you HuggingFace account, you can get your free API token in your settings page: https://huggingface.co/settings/tokens",
+    #     type="password",
+    # )
+    #
+    # # Adding the HuggingFace API inference URL.
+    # API_URL = "https://api-inference.huggingface.co/models/valhalla/distilbart-mnli-12-3"
+    #
+    # # Now, let's create a Python dictionary to store the API headers.
+    # headers = {"Authorization": f"Bearer {API_KEY}"}
 
 
     st.sidebar.markdown("---")
@@ -77,6 +84,25 @@ def main(model,tokenizer):
     """
     )
 
+    generator = pipeline('text-generation', model="hanbin/MaMaL-Com")
+
+
+
+    # model, tokenizer = load_model("hanbin/MaMaL-Gen")
+    st.write("### 输入：")
+    input = st.text_area("",height=200)
+    output = generator(input)
+    # code = '''def hello():
+    #     print("Hello, Streamlit!")'''
+    if st.button('补全'):
+        st.write("### 输出：")
+        st.code(output, language='python')
+    else:
+        st.write('')
+
+
+
+
 if __name__ == '__main__':
-    model, tokenizer = load_model("hanbin/MaMaL-Gen")
-    main(model, tokenizer)
+
+    main()