v1 complete

.gitignore

@@ -171,4 +171,5 @@ tags
 # local
 *.ipynb
 docs/
-repos/
+repos/
+milvus.db

app.py

@@ -3,8 +3,28 @@ import os
 import json
 import subprocess
 import tempfile
+import dotenv
 import shutil
 from pathlib import Path
+from string import Template
+from pymilvus import MilvusClient, model
+
+assert dotenv.load_dotenv()
+
+template = Template("""\
+---
+File: $file_path
+---
+                    
+$file_content""")
+
+client = MilvusClient("milvus.db")
+embedding_fn = model.dense.OpenAIEmbeddingFunction(
+    model_name='text-embedding-3-small', # Specify the model name
+    api_key=os.environ.get('OPENAI_API_KEY'), # Provide your OpenAI API key
+    dimensions=1536 # Set the embedding dimensionality
+)
+
 
 def list_huggingface_resources_names() -> list[str]:
     """List all the names of the libraries, services, and other resources available within the HuggingFace ecosystem.
@@ -15,9 +35,38 @@ def list_huggingface_resources_names() -> list[str]:
     with open('repos_config.json', 'r') as f:
         repos = json.load(f)
 
+    print([repo['title'] for repo in repos])
+
     return [repo['title'] for repo in repos]
 
 
+def get_huggingface_documentation(topic: str, resource_names: list[str]) -> str:
+    """Get the documentation for the given topic and resource names.
+    
+    Args:
+        topic: Focus the docs on a specific topic (e.g. "Anthropic Provider Chat UI", "LoRA methods PEFT" or "TGI on Intel GPUs")
+        resource_names: A list of relevant resource names to the topic
+        
+    Returns:
+        A string of documentation for the given topic and resource names
+    """
+    print(resource_names)
+    query_vectors = embedding_fn.encode_queries([topic])
+    res = client.search(collection_name="hf_docs", data=query_vectors, limit=3, output_fields=["text", "file_path"])
+    print(res)
+
+    docs_paths = [res[0][i]['file_path'] for i in range(len(res[0]))]
+    print(docs_paths)
+
+    documentation = ""
+    for path in docs_paths:
+        with open(path, 'r') as f:
+            content = f.read()
+            documentation += template.substitute(file_path=path.replace('docs/', ''), file_content=content) + "\n\n"
+
+    print(documentation.strip())
+    return documentation.strip()
+
 list_resources_demo = gr.Interface(
     fn=list_huggingface_resources_names,
     inputs=[],
@@ -26,11 +75,17 @@ list_resources_demo = gr.Interface(
     description="Explore the names of the libraries, services, and other resources available within the HuggingFace ecosystem"
 )
 
+get_docs_demo = gr.Interface(
+    fn=get_huggingface_documentation,
+    inputs=["text", "json"],
+    outputs="text",
+)
+
 # Create tabbed interface
 demo = gr.TabbedInterface(
-    [list_resources_demo],
-    ["List Resources"],
+    [list_resources_demo, get_docs_demo],
+    ["List Resources", "Get Documentation"],
     title="HuggingFace Ecosystem Documentation Explorer",
-)
+)   
 
 demo.launch(mcp_server=True)

make_docs.py

@@ -93,7 +93,7 @@ def save_section_to_disk(section: Dict, file_path: Path, raw_docs_path: Path):
             shutil.copy(local_path, file_path / f"{title}{local_path.suffix}")
         
         except Exception as e:
-            # TODO: Handle symlinks, missing files, and other edge cases
+            # TODO: Not many cases, but handle symlinks, missing files, and other edge cases
             pass
 
 
@@ -128,4 +128,5 @@ if __name__ == "__main__":
     with open("repos_config.json", "r") as f:
         repos = json.load(f)
 
+    shutil.rmtree(args.docs_dir)
     make_docs(repos, args)

make_rag_db.py

@@ -0,0 +1,62 @@
+import os
+import argparse
+from typing import Dict
+import dotenv
+from pathlib import Path
+from tqdm import tqdm
+from pymilvus import MilvusClient, model
+
+assert dotenv.load_dotenv()
+
+
+def create_collection(client: MilvusClient, collection_name: str, dimension: int):
+    if client.has_collection(collection_name=collection_name):
+        client.drop_collection(collection_name=collection_name)
+
+    client.create_collection(
+        collection_name=collection_name,
+        dimension=dimension,
+    )
+
+def main(args: Dict):
+    client = MilvusClient("milvus.db")
+
+    embedding_fn = model.dense.OpenAIEmbeddingFunction(
+        model_name=args.model_name,
+        api_key=os.environ.get('OPENAI_API_KEY'),
+        dimensions=args.dimension
+    )
+
+    create_collection(client, args.collection_name, args.dimension)
+
+    docs = Path(args.docs_dir)
+    md_file_paths  = list(docs.rglob('*.md'))
+    mdx_file_paths = list(docs.rglob('*.mdx'))
+    all_file_paths = md_file_paths + mdx_file_paths
+    
+    docs, payloads = [], []
+    for file in tqdm(all_file_paths):
+        embed_string = str(file).replace('docs/', '').replace('.mdx', '').replace('.md', '').replace('/', ' ')
+
+        docs.append(embed_string)
+        payloads.append({'file_path': str(file)})
+
+    vectors = embedding_fn.encode_documents(docs)
+
+    data = [
+        {"id": i, "vector": vectors[i], "text": docs[i], **payloads[i]}
+        for i in range(len(vectors))
+    ]
+
+    response = client.insert(collection_name=args.collection_name, data=data)
+    print(f"Inserted {response['insert_count']} vectors into collection {args.collection_name}")
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--collection_name", type=str, default="hf_docs")
+    parser.add_argument("--model_name", type=str, default="text-embedding-3-small")
+    parser.add_argument("--dimension", type=int, default=1536)
+    parser.add_argument("--docs_dir", type=str, default="docs")
+    args = parser.parse_args()
+
+    main(args)

postBuild

@@ -0,0 +1,2 @@
+python3 make_docs.py
+python3 make_rag_db.py

requirements.txt

@@ -0,0 +1,5 @@
+pymilvus==2.5.10
+pymilvus_model==0.3.2
+python-dotenv==1.1.0
+PyYAML==6.0.2
+tqdm==4.65.0