update

Dockerfile

@@ -1,4 +1,4 @@
-FROM python:3.9.2-slim
+FROM --platform=linux/amd64 python:3.9.2-slim
 RUN apt-get update && apt-get install -y gcc g++
 WORKDIR /app
 

README.md

@@ -1,6 +1,6 @@
 ---
-title: Streamlit Docker Template
-emoji: 📉
+title: Polos Demo
+emoji: 🌟
 colorFrom: blue
 colorTo: green
 sdk: docker
@@ -8,8 +8,14 @@ app_port: 8501
 pinned: false
 ---
 
-## 🧠 Streamlit Docker Template 🔎
+## Get Started on M1 Mac
+```bash
+git submodule update --init --recursive
+docker build . -t polos_demo
+docker run -it -d -v `pwd`:/workspace -p 8080:8080 --platform linux/amd64 polos_demo
+docker exec -it $process_id bash
 
-Streamlit Docker Template is a template for creating a Streamlit app with Docker and Hugging Face Spaces.
-
-Code from https://docs.streamlit.io/library/get-started/create-an-app
+root@28cb354f7609:~# sh install.sh
+root@28cb354f7609:~# poetry run python test.py
+root@28cb354f7609:~# poetry run streamlit run test.py --server.port 8080
+```

app.py

@@ -2,7 +2,8 @@ import streamlit as st
 from PIL import Image
 from polos.models import download_model, load_checkpoint
 
-@st.cache(allow_output_mutation=True)
+# モデルのロード
+@st.cache_resource()
 def load_model():
     model_path = download_model("polos")
     model = load_checkpoint(model_path)
@@ -10,29 +11,67 @@ def load_model():
 
 model = load_model()
 
-default_image = Image.open("test.jpg").convert("RGB")
-default_refs = [
-    "there is a dog sitting on a couch with a person reaching out",
-    "a dog laying on a couch with a person",
-    'a dog is laying on a couch with a person'
-]
-
-data = [
-    {
-        "img": default_image,
-        "mt": "",
-        "refs": default_refs
-    }
-]
-
 # Streamlitインターフェースの設定
 st.title('Polos Demo')
 
+# セッションステートの初期化
+if 'image' not in st.session_state:
+    st.session_state.image = None
+if 'user_input' not in st.session_state:
+    st.session_state.user_input = ''
+if 'user_refs' not in st.session_state:
+    st.session_state.user_refs = [
+        "there is a dog sitting on a couch with a person reaching out",
+        "a dog laying on a couch with a person",
+        'a dog is laying on a couch with a person'
+    ]
+if 'score' not in st.session_state:
+    st.session_state.score = None
+
+# デフォルト画像の取得
+@st.cache_resource()
+def get_default_image():
+    try:
+        return Image.open("test.jpg").convert("RGB")
+    except FileNotFoundError:
+        return Image.new('RGB', (200, 200), color = 'gray')  # デフォルト画像が見つからない場合の代替画像
+
+default_image = get_default_image()
+
+# 画像アップロードのためのウィジェット
+uploaded_image = st.file_uploader("Upload your image:", type=["jpg", "jpeg", "png"])
+if uploaded_image is not None:
+    st.session_state.image = Image.open(uploaded_image).convert("RGB")
+elif st.session_state.image is None:
+    st.session_state.image = default_image
+
+# 常に画像を表示
+st.image(st.session_state.image, caption="Displayed Image", use_column_width=True)
+
+# 参照文の入力フィールド
+user_refs = st.text_area("Enter reference sentences (separate each by a newline):", "\n".join(st.session_state.user_refs))
+st.session_state.user_refs = user_refs.split("\n")
+
 # ユーザー入力のテキストフィールド
-user_input = st.text_input("Enter the input sentence:", '')
+user_input = st.text_input("Enter the input sentence:", value=st.session_state.user_input)
+st.session_state.user_input = user_input
+
+# Computeボタン
+if st.button('Compute'):
+    # データの準備
+    data = [
+        {
+            "img": st.session_state.image,
+            "mt": st.session_state.user_input,
+            "refs": st.session_state.user_refs
+        }
+    ]
+
+    # モデル予測
+    if st.session_state.user_input:
+        _, scores = model.predict(data, batch_size=1, cuda=False)
+        st.session_state.score = scores[0]
 
-# 入力がある場合、モデルを使用してスコアを計算
-if user_input:
-    data[0]['mt'] = user_input
-    _, scores = model.predict(data, batch_size=1, cuda=False)
-    st.write("Score:", scores)
+# スコアの表示
+if st.session_state.score is not None:
+    st.metric(label="Score", value=f"{st.session_state.score:.5f}")

polos/models/__init__.py

@@ -19,13 +19,10 @@ str2model = {
 }
 
 def get_cache_folder():
-    if "HOME" in os.environ:
-        cache_directory = os.environ["HOME"] + "/.cache/torch/yuigawada/"
-        if not os.path.exists(cache_directory):
-            os.makedirs(cache_directory)
-        return cache_directory
-    else:
-        raise Exception("HOME environment variable is not defined.")
+    cache_directory = "./.cache/"
+    if not os.path.exists(cache_directory):
+        os.makedirs(cache_directory)
+    return cache_directory
 
 def download_model(model: str, saving_directory: str = None) -> ModelBase:
     """Function that loads pretrained models from AWS.

polos/models/encoders/xlmr.py

@@ -29,10 +29,7 @@ XLMR_LARGE_V0_MODEL_NAME = "xlmr.large.v0/model.pt"
 XLMR_BASE_V0_URL = "https://dl.fbaipublicfiles.com/fairseq/models/xlmr.base.v0.tar.gz"
 XLMR_BASE_V0_MODEL_NAME = "xlmr.base.v0/model.pt"
 
-if "HOME" in os.environ:
-    saving_directory = os.environ["HOME"] + "/.cache/torch/yuigawada/"
-else:
-    raise Exception("HOME environment variable is not defined.")
+saving_directory = "./.cache/"
 
 
 class XLMREncoder(Encoder):