build a github similarity score retriever - no streamlit integration yet

.gitignore

@@ -0,0 +1,3 @@
+.env
+.venv/
+lib/

README.md

@@ -11,3 +11,9 @@ license: mit
 ---
 
 
+## Local Dev
+
+`python -m venv .venv`
+`source .venv/bin/activate/`
+`pip install -r requirements.txt`
+`streamlit run app.py`

app.py

@@ -5,6 +5,7 @@ import pickle
 import torch
 import io
 from langchain.vectorstores import FAISS
+import json
 
 class CPU_Unpickler(pickle.Unpickler):
     def find_class(self, module, name):
@@ -24,8 +25,29 @@ def get_hugging_face_model():
 def get_db():
   with open("codesearchdb.pickle", "rb") as f:
     db = CPU_Unpickler(f).load()
+  print("Loaded db")
+  # save_as_json(db, "codesearchdb.json")  # Save as JSON
   return db
 
+def save_as_json(data, filename):
+    # Convert the data to a JSON serializable format
+    serializable_data = data_to_serializable(data)
+    with open(filename, "w") as json_file:
+        json.dump(serializable_data, json_file)
+
+def data_to_serializable(data):
+    if isinstance(data, dict):
+        return {k: data_to_serializable(v) for k, v in data.items() if not callable(v) and not isinstance(v, type)}
+    elif isinstance(data, list):
+        return [data_to_serializable(item) for item in data]
+    elif isinstance(data, (str, int, float, bool)) or data is None:
+        return data
+    elif hasattr(data, '__dict__'):
+        return data_to_serializable(data.__dict__)
+    elif hasattr(data, '__slots__'):
+        return {slot: data_to_serializable(getattr(data, slot)) for slot in data.__slots__}
+    else:
+        return str(data)  # Convert any other types to string
 
 def get_similar_links(query, db, embeddings):
   embedding_vector = embeddings.embed_query(query)
@@ -45,6 +67,7 @@ def get_similar_links(query, db, embeddings):
 
 embedding_vector = get_hugging_face_model()
 db = FAISS.load_local("code_sim_index", embedding_vector, allow_dangerous_deserialization=True)
+save_as_json(db, "code_sim_index.json")  # Save as JSON
 
 st.title("Find Similar Code")
 text_input = st.text_area("Enter a Code Example", value =
@@ -73,3 +96,4 @@ if button:
 else:
   st.info("Please Input Valid Text")
 
+# get_db()

github.py

@@ -0,0 +1,67 @@
+import os
+from dotenv import load_dotenv
+from langchain.document_loaders import GithubFileLoader
+# from langchain.embeddings import HuggingFaceEmbeddings
+from langchain_huggingface import HuggingFaceEmbeddings
+from langchain_community.vectorstores import FAISS
+from langchain_text_splitters import CharacterTextSplitter
+
+load_dotenv()
+
+#get the GITHUB_ACCESS_TOKEN from the .env file
+GITHUB_ACCESS_TOKEN = os.getenv("GITHUB_ACCESS_TOKEN")
+USER = "heaversm"
+REPO = "gdrive-docker"
+GITHUB_BASE_URL = "https://github.com/"
+
+
+def get_similar_files(query, db, embeddings):
+  # embedding_vector = embeddings.embed_query(query)
+  # docs_and_scores = db.similarity_search_by_vector(embedding_vector, k = 10)
+  docs_and_scores = db.similarity_search_with_score(query)
+  return docs_and_scores
+
+def get_hugging_face_model():
+  model_name = "mchochlov/codebert-base-cd-ft"
+  hf = HuggingFaceEmbeddings(model_name=model_name)
+  return hf
+
+loader = GithubFileLoader(
+    #repo is USER/REPO
+    repo=f"{USER}/{REPO}",
+    access_token=GITHUB_ACCESS_TOKEN,
+    github_api_url="https://api.github.com",
+    file_filter=lambda file_path: file_path.endswith(
+        (".py", ".ts")
+    ),  # load all python and typescript files
+)
+documents = loader.load()
+text_splitter = CharacterTextSplitter(chunk_size=1000, chunk_overlap=0)
+docs = text_splitter.split_documents(documents)
+embedding_vector = get_hugging_face_model()
+db = FAISS.from_documents(docs, embedding_vector)
+model_name = "mchochlov/codebert-base-cd-ft"
+
+query = """
+def create_app():
+    app = connexion.FlaskApp(__name__, specification_dir="../.openapi")
+    app.add_api(
+        API_VERSION, resolver=connexion.resolver.RelativeResolver("provider.app")
+    )
+"""
+results_with_scores = get_similar_files(query, db, embedding_vector)
+print ("retrieved!!!")
+print(f"Number of results: {len(results_with_scores)}")
+
+# score is a distance score, the lower the better
+for doc, score in results_with_scores:
+    print(f"Metadata: {doc.metadata}, Score: {score}")
+
+top_file_path = results_with_scores[0][0].metadata['path']
+top_file_content = results_with_scores[0][0].page_content
+top_file_score = results_with_scores[0][1]
+top_file_link = f"{GITHUB_BASE_URL}{USER}/{REPO}/blob/main/{top_file_path}"
+
+print(f"Top file link: {top_file_link}")
+
+

requirements.txt

@@ -3,4 +3,6 @@ sentence-transformers
 bs4
 faiss-cpu
 altair==4.0
-langchain-community
+langchain-community
+streamlit
+python-dotenv

search-pickle.py

@@ -0,0 +1,99 @@
+import streamlit as st
+from bs4 import BeautifulSoup
+from langchain.embeddings import HuggingFaceEmbeddings
+import pickle
+import torch
+import io
+from langchain.vectorstores import FAISS
+import json
+
+class CPU_Unpickler(pickle.Unpickler):
+    def find_class(self, module, name):
+        if module == 'torch.storage' and name == '_load_from_bytes':
+            return lambda b: torch.load(io.BytesIO(b), map_location='cpu')
+        else: return super().find_class(module, name)
+
+
+@st.cache_resource
+def get_hugging_face_model():
+  model_name = "mchochlov/codebert-base-cd-ft"
+  hf = HuggingFaceEmbeddings(model_name=model_name)
+  return hf
+
+
+@st.cache_resource
+def get_db():
+  with open("codesearchdb.pickle", "rb") as f:
+    db = CPU_Unpickler(f).load()
+  print("Loaded db")
+  # save_as_json(db, "codesearchdb.json")  # Save as JSON
+  return db
+
+def save_as_json(data, filename):
+    # Convert the data to a JSON serializable format
+    serializable_data = data_to_serializable(data)
+    with open(filename, "w") as json_file:
+        json.dump(serializable_data, json_file)
+
+def data_to_serializable(data):
+    if isinstance(data, dict):
+        return {k: data_to_serializable(v) for k, v in data.items() if not callable(v) and not isinstance(v, type)}
+    elif isinstance(data, list):
+        return [data_to_serializable(item) for item in data]
+    elif isinstance(data, (str, int, float, bool)) or data is None:
+        return data
+    elif hasattr(data, '__dict__'):
+        return data_to_serializable(data.__dict__)
+    elif hasattr(data, '__slots__'):
+        return {slot: data_to_serializable(getattr(data, slot)) for slot in data.__slots__}
+    else:
+        return str(data)  # Convert any other types to string
+
+def get_similar_links(query, db, embeddings):
+  embedding_vector = embeddings.embed_query(query)
+  docs_and_scores = db.similarity_search_by_vector(embedding_vector, k = 10)
+  hrefs = []
+  for docs in docs_and_scores:
+    html_doc = docs.page_content
+    soup = BeautifulSoup(html_doc, 'html.parser')
+    href = [a['href'] for a in soup.find_all('a', href=True)]
+    hrefs.append(href)
+  links = []
+  for href_list in hrefs:
+      for link in href_list:
+        links.append(link)
+  return links
+
+
+embedding_vector = get_hugging_face_model()
+db = FAISS.load_local("code_sim_index", embedding_vector, allow_dangerous_deserialization=True)
+save_as_json(db, "code_sim_index.json")  # Save as JSON
+
+st.title("Find Similar Code")
+text_input = st.text_area("Enter a Code Example", value =
+"""
+class Solution:
+    def subsets(self, nums: List[int]) -> List[List[int]]:
+        outputs = []
+        def backtrack(k, index, subSet):
+            if index == k:
+                outputs.append(subSet[:])
+                return
+            for i in range(index, len(nums)):
+                backtrack(k, i + 1, subSet + [nums[i]])
+        for j in range(len(nums) + 1):
+            backtrack(j, 0, [])
+        return outputs
+""", height = 330
+)
+button = st.button("Find Similar Questions")
+if button:
+  query = text_input
+  answer = get_similar_links(query, db, embedding_vector)
+  for link in set(answer):
+      st.write(link)
+
+else:
+  st.info("Please Input Valid Text")
+
+# get_db()

unpickle.py

@@ -0,0 +1,19 @@
+import pickle
+
+# Define the path to the pickle file
+pickle_file_path = 'codesearchdb.pickle'
+
+# Load the pickle file
+with open(pickle_file_path, 'rb') as file:
+    data = pickle.load(file)
+
+
+
+# Save the contents to a new file (for example, a JSON file)
+import json
+
+json_file_path = 'codesearchdb.json'
+with open(json_file_path, 'w') as json_file:
+    json.dump(data, json_file, indent=4)
+
+print(f"Contents have been saved to {json_file_path}")