init space

config.py

@@ -0,0 +1,29 @@
+from os import getenv
+from dotenv import load_dotenv
+
+from pydantic import BaseSettings
+
+
+class Config(BaseSettings):
+    """Configuration for the application."""
+    # Pinecone
+    pinecone_api_key: str
+    pinecone_env: str
+    pinecone_index: str
+    # LLM
+    embedding_dim: int
+    embedding_version: str
+    embedding_dir: str
+    model_name: str
+
+
+load_dotenv()
+config = Config(
+    pinecone_api_key=getenv("PINECONE_API_KEY"),
+    pinecone_env=getenv("PINECONE_ENV"),
+    pinecone_index=getenv("PINECONE_INDEX"),
+    embedding_dim=int(getenv("EMBEDDING_DIM")),
+    embedding_version=getenv("EMBEDDING_VERSION"),
+    embedding_dir=getenv("EMBEDDING_DIR"),
+    model_name=getenv("MODEL_NAME"),
+)

main.py

@@ -0,0 +1,64 @@
+import pinecone
+import requests
+import streamlit as st
+import torch
+
+from transformers import AutoTokenizer, AutoModel
+
+from config import config
+    
+    
+def search(text: str, k: int = 5):
+    """Get the k closest articles to the text."""
+    embeds = _get_embeddings(text)
+    
+    r = requests.post(
+        f"https://{config.pinecone_index}-5b18b87.svc.{config.pinecone_env}.pinecone.io/query",
+        headers={
+            "Api-Key": config.pinecone_api_key,
+            "accept": "application/json",
+            "content-type": "application/json",
+        },
+        json={
+            "vector": embeds,
+            "top_k": k,
+            "includeMetadata": True,
+            "includeValues": False,
+        },
+    )
+    
+    if r.status_code == 200:
+        return r.json()
+    else:
+        raise Exception(f"Error: {r.status_code} - {r.text}")
+
+    
+def _get_embeddings(text: str):
+    inputs_ids = st.session_state.tokenizer(text, return_tensors="pt", padding=True, truncation=True)
+    
+    with torch.no_grad():
+        last_hidden_states = st.session_state.model(**inputs_ids)[0]
+    
+    return last_hidden_states.mean(dim=1).squeeze().tolist()
+
+
+st.title("PubMed Embeddings")
+st.subheader("Search for a PubMed article and get its id.")
+    
+text = st.text_input("Search for a PubMed article", "Epidemiology of COVID-19")
+
+with st.spinner("Loading Embedding Model..."):
+    pinecone.init(api_key=config.pinecone_api_key, env=config.pinecone_env)
+    if "index" not in st.session_state:
+        st.session_state.index = pinecone.Index(config.pinecone_index)
+    if "tokenizer" not in st.session_state:          
+        st.session_state.tokenizer = AutoTokenizer.from_pretrained(config.model_name)
+    if "model" not in st.session_state:
+        st.session_state.model = AutoModel.from_pretrained(config.model_name)
+
+if st.button("Search"):
+    with st.spinner("Searching..."):
+        results = st.session_state.embeds_handler.search(text)
+
+    for res in results["matches"]:
+        st.write(f"{res['id']} - confidence: {res['score']:.2f}")

requirements.txt

@@ -0,0 +1,7 @@
+pinecone-client>=2.1.0
+python-dotenv>=0.21.1
+pydantic>=1.10.4
+requests>=2.26.0
+streamlit>=1.17.0
+transformers>=4.26.0
+torch>=1.12.0