project init

.flake8

@@ -0,0 +1,7 @@
+[flake8]
+ignore = E203, E266, E501, W503, B904
+# line length is intentionally set to 80 here because black uses Bugbear
+# See https://black.readthedocs.io/en/stable/the_black_code_style/current_style.html#line-length for more details
+max-line-length = 80
+max-complexity = 18
+select = B, C, E, F, W, T4, B9

.pre-commit-config.yaml

@@ -0,0 +1,31 @@
+repos:
+  - repo: https://github.com/pre-commit/pre-commit-hooks
+    rev: v4.2.0
+    hooks:
+      - id: trailing-whitespace
+        args: [--markdown-linebreak-ext=md]
+      - id: end-of-file-fixer
+
+  - repo: https://github.com/asottile/pyupgrade
+    rev: v2.32.1
+    hooks:
+      - id: pyupgrade
+        args: [--py39-plus]
+
+  - repo: https://github.com/pycqa/isort
+    rev: 5.10.1
+    hooks:
+      - id: isort
+        args: [--profile=black]
+
+  - repo: https://github.com/psf/black
+    rev: 22.3.0
+    hooks:
+      - id: black
+
+  - repo: https://github.com/PyCQA/flake8
+    rev: 4.0.1
+    hooks:
+      - id: flake8
+        additional_dependencies:
+          - flake8-bugbear

app.py

@@ -0,0 +1,28 @@
+import pickle
+
+import streamlit as st
+import numpy as np
+import pandas as pd
+from transformers import VisionTextDualEncoderModel, VisionTextDualEncoderProcessor
+from sentence_transformers.util import semantic_search
+
+
+st.title("VitB32 Bert Ko Small Clip Test")
+
+model = VisionTextDualEncoderModel.from_pretrained("Bingsu/vitB32_bert_ko_small_clip")
+processor = VisionTextDualEncoderProcessor.from_pretrained(
+    "Bingsu/vitB32_bert_ko_small_clip"
+)
+
+info = pd.read_csv("info.csv")
+with open("img_id.pkl", "rb") as f:
+    img_id = pickle.load(f)
+img_emb = np.load("img_emb.npy")
+
+text = st.text_input("Input Text")
+tokens = processor(text=text, return_tensors="pt")
+text_emb = model.get_text_features(**tokens)
+
+result = semantic_search(text_emb, img_emb, top_k=6)
+
+st.write(result)

pyproject.toml

@@ -0,0 +1,29 @@
+[tool.poetry]
+name = "vitb32_bert_ko_small_clip_test"
+version = "0.1.0"
+description = ""
+authors = ["Bingsu <ks2515@naver.com>"]
+license = "MIT"
+
+[tool.poetry.dependencies]
+python = "^3.9"
+torch = "^1.11.0"
+transformers = "^4.19.2"
+streamlit = "^1.10.0"
+pandas = "^1.4.2"
+sentence-transformers = "^2.2.0"
+
+[tool.poetry.dev-dependencies]
+black = "^22.3.0"
+isort = "^5.10.1"
+mypy = "^0.961"
+flake8 = "^4.0.1"
+flake8-bugbear = "^22.4.25"
+pre-commit = "^2.19.0"
+
+[build-system]
+requires = ["poetry-core>=1.0.0"]
+build-backend = "poetry.core.masonry.api"
+
+[tool.isort]
+profile="black"

requirements.txt

@@ -0,0 +1,4 @@
+torch
+transformers
+pandas
+sentence-transformers