| import gradio as gr |
| from sentence_transformers import SentenceTransformer |
| import numpy as np |
|
|
| MODEL_NAME = "Matthieufromparis/bge-small-code-search-v1" |
|
|
| model = SentenceTransformer(MODEL_NAME) |
|
|
| CODE_SNIPPETS = [ |
| "def parse_json_config(filepath):\n with open(filepath) as f:\n return json.load(f)", |
| "def sort_by_key(items, key, reverse=False):\n return sorted(items, key=lambda x: x.get(key, ''), reverse=reverse)", |
| "def authenticate_user(token):\n payload = jwt.decode(token, SECRET, algorithms=['HS256'])\n return payload.get('user_id')", |
| "def fetch_api_data(url, headers=None):\n resp = requests.get(url, headers=headers, timeout=10)\n resp.raise_for_status()\n return resp.json()", |
| "def validate_email(email):\n pattern = r'^[a-zA-Z0-9._%+-]+@[a-zA-Z0-9.-]+\\.[a-zA-Z]{2,}$'\n return bool(re.match(pattern, email))", |
| "def connect_database(host, port, user, password, dbname):\n conn = psycopg2.connect(host=host, port=port, user=user, password=password, dbname=dbname)\n return conn", |
| "def retry_on_failure(func, max_retries=3, delay=1):\n for attempt in range(max_retries):\n try:\n return func()\n except Exception as e:\n if attempt == max_retries - 1:\n raise\n time.sleep(delay)", |
| "def generate_random_string(length=12):\n chars = string.ascii_letters + string.digits\n return ''.join(random.choice(chars) for _ in range(length))", |
| "def log_error(message, level='ERROR'):\n logger = logging.getLogger(__name__)\n getattr(logger, level.lower())(message)", |
| "def chunk_list(items, chunk_size=100):\n for i in range(0, len(items), chunk_size):\n yield items[i:i + chunk_size]", |
| ] |
|
|
| code_embeddings = model.encode(CODE_SNIPPETS, normalize_embeddings=True) |
|
|
| def search_code(query: str, top_k: int = 5): |
| if not query.strip(): |
| return "Enter a query to search code snippets." |
| query_embedding = model.encode(query, normalize_embeddings=True) |
| scores = np.dot(query_embedding, code_embeddings.T) |
| top_indices = np.argsort(scores)[-top_k:][::-1] |
| results = [] |
| for idx in top_indices: |
| score = float(scores[idx]) |
| if score < 0.3: |
| continue |
| results.append(f"### Score: {score:.3f}\n```python\n{CODE_SNIPPETS[idx]}\n```\n---") |
| if not results: |
| return "No matching code snippets found. Try a different query." |
| return "\n".join(results) |
|
|
| def compare_embeddings(text1, text2): |
| if not text1.strip() or not text2.strip(): |
| return "Enter both texts to compare." |
| emb1 = model.encode(text1, normalize_embeddings=True) |
| emb2 = model.encode(text2, normalize_embeddings=True) |
| similarity = float(np.dot(emb1, emb2)) |
| if similarity > 0.8: level = "🟢 Very Similar" |
| elif similarity > 0.6: level = "🟡 Somewhat Similar" |
| elif similarity > 0.4: level = "🟠 Slightly Similar" |
| else: level = "🔴 Not Similar" |
| return f"## {level}\n**Cosine Similarity: {similarity:.4f}**\n\n*384-dimensional embeddings*" |
|
|
| with gr.Blocks(title="Code Embeddings Demo — Matthieu.AI", theme=gr.themes.Soft()) as demo: |
| gr.Markdown("""# 🔍 bge-small-code-search-v1\n### Semantic Code Search Demo\nSearch code by describing what it does. Model by [Matthieu.AI](https://huggingface.co/Matthieufromparis)""") |
| with gr.Tab("Code Search"): |
| query_input = gr.Textbox(label="Search Query", placeholder="e.g., 'parse a JSON config file'", lines=2) |
| top_k = gr.Slider(minimum=1, maximum=10, value=5, step=1, label="Results") |
| search_button = gr.Button("Search Code", variant="primary") |
| search_output = gr.Markdown(label="Results") |
| search_button.click(fn=search_code, inputs=[query_input, top_k], outputs=search_output) |
| gr.Examples(examples=[["read a JSON config file"], ["validate an email address"], ["connect to a database"], ["retry a function if it fails"]], inputs=[query_input]) |
| with gr.Tab("Compare Texts"): |
| text1 = gr.Textbox(label="Text 1", placeholder="Enter natural language or code...", lines=3) |
| text2 = gr.Textbox(label="Text 2", placeholder="Enter another text to compare...", lines=3) |
| compare_button = gr.Button("Compare", variant="primary") |
| compare_output = gr.Markdown(label="Similarity") |
| compare_button.click(fn=compare_embeddings, inputs=[text1, text2], outputs=compare_output) |
|
|
| if __name__ == "__main__": |
| demo.launch() |
|
|
