integrated hugging face dataset

Author: Gerome-Elassaad

.github/workflows/changelog.yml

@@ -0,0 +1,31 @@
+name: Release
+
+on:
+  push:
+    branches:
+      - main
+
+jobs:
+  release:
+    runs-on: ubuntu-latest
+    permissions:
+      contents: write
+      
+    steps:
+      - name: Checkout code
+        uses: actions/checkout@v4
+
+      - name: Bump version and push tag
+        id: tag_version
+        uses: mathieudutour/github-tag-action@v6.2
+        with:
+          github_token: ${{ secrets.GITHUB_TOKEN }}
+
+      - name: Create a GitHub release
+        uses: actions/create-release@v1
+        with:
+          tag_name: ${{ steps.tag_version.outputs.new_tag }}
+          release_name: "Release ${{ steps.tag_version.outputs.new_tag }}"
+          body: ${{ steps.tag_version.outputs.changelog }}
+        env:
+          GITHUB_TOKEN: ${{ secrets.GITHUB_TOKEN }}

.github/workflows/release.yml

@@ -36,3 +36,5 @@ yarn-error.log*
 # typescript
 *.tsbuildinfo
 next-env.d.ts
+
+gcp.json

.gitignore

@@ -3,6 +3,7 @@ import { getModelClient } from '@/lib/models'
 import { LLMModel, LLMModelConfig } from '@/lib/models'
 import { toPrompt } from '@/lib/prompt'
 import ratelimit from '@/lib/ratelimit'
+import { retrieveRelevantCode } from '@/lib/retrieval'
 import { fragmentSchema as schema } from '@/lib/schema'
 import { Templates } from '@/lib/templates'
 import { streamObject, LanguageModel, CoreMessage } from 'ai'
@@ -53,19 +54,34 @@ export async function POST(req: Request) {
   }
 
   console.log('userID', userID)
-  console.log('teamID', teamID)
-  // console.log('template', template)
+  console.log('template', template)
   console.log('model', model)
-  // console.log('config', config)
+  console.log('config', config)
 
   const { model: modelNameString, apiKey: modelApiKey, ...modelParams } = config
   const modelClient = getModelClient(model, config)
 
   try {
+    const lastUserMessage = messages[messages.length - 1]?.content;
+    let retrievedContext = '';
+    if (typeof lastUserMessage === 'string') {
+      const relevantCode = await retrieveRelevantCode(lastUserMessage);
+      if (relevantCode.length > 0) {
+        retrievedContext = `
+---
+Here is some relevant code from our knowledge base that might help:
+${relevantCode.join('\n---\n')}
+---
+`;
+      }
+    }
+
+    const systemPrompt = toPrompt(template) + retrievedContext;
+
     const stream = await streamObject({
       model: modelClient as LanguageModel,
       schema,
-      system: toPrompt(template),
+      system: systemPrompt,
       messages,
       maxRetries: 0, // do not retry on errors
       ...modelParams,
@@ -116,4 +132,4 @@ export async function POST(req: Request) {
       },
     )
   }
-}
+}

app/api/chat/route.ts

@@ -0,0 +1,102 @@
+import { NextResponse } from 'next/server';
+import { pipeline } from '@xenova/transformers';
+import { createClient } from '@supabase/supabase-js';
+
+// Configuration
+const DATASET = 'bigcode/the-stack';
+const EMBEDDING_MODEL = 'Xenova/all-MiniLM-L6-v2';
+const HUGGING_FACE_API_URL = 'https://huggingface.co/api/datasets';
+
+// Initialize Supabase client with the service role key for admin access
+const supabase = createClient(
+  process.env.NEXT_PUBLIC_SUPABASE_URL!,
+  process.env.SUPABASE_SERVICE_ROLE_KEY!
+);
+
+// Function to process and embed a single file's content
+async function processAndEmbedFile(content: string, embeddingPipeline: any) {
+  // 1. Chunk the content
+  const chunks = content.split('\n').filter(line => line.trim().length > 10);
+
+  if (chunks.length === 0) {
+    return;
+  }
+
+  // 2. Generate embeddings for each chunk
+  const embeddings = await embeddingPipeline(chunks, {
+    pooling: 'mean',
+    normalize: true,
+  });
+
+  // 3. Prepare data for Supabase
+  const dataToInsert = chunks.map((chunk, i) => ({
+    content: chunk,
+    embedding: Array.from(embeddings.data.slice(i * 384, (i + 1) * 384)),
+  }));
+
+  // 4. Upsert into Supabase
+  const { error } = await supabase.from('code_embeddings').insert(dataToInsert);
+  if (error) {
+    console.error('Supabase insert error:', error);
+  }
+}
+
+// The main background job logic
+async function importDataset(subset: string) {
+  try {
+    console.log(`Starting dataset import for subset: ${subset}...`);
+    const embeddingPipeline = await pipeline('feature-extraction', EMBEDDING_MODEL);
+
+    // 1. Fetch the list of files in the dataset subset
+    const repoInfoUrl = `${HUGGING_FACE_API_URL}/${DATASET}/tree/main/data/${subset}`;
+    const repoInfoResponse = await fetch(repoInfoUrl);
+    if (!repoInfoResponse.ok) {
+      throw new Error(`Failed to fetch repo info: ${repoInfoResponse.statusText}`);
+    }
+    const files: { path: string; type: string }[] = await repoInfoResponse.json();
+
+    // 2. Process each file
+    for (const file of files) {
+      if (file.type === 'file') {
+        console.log(`Processing file: ${file.path}`);
+        
+        // Construct the direct download URL
+        const downloadUrl = `https://huggingface.co/datasets/${DATASET}/resolve/main/${file.path}`;
+        const fileResponse = await fetch(downloadUrl);
+        
+        if (fileResponse.ok) {
+          const content = await fileResponse.text();
+          await processAndEmbedFile(content, embeddingPipeline);
+        } else {
+          console.warn(`Could not download file: ${file.path}`);
+        }
+      }
+    }
+
+    console.log(`Dataset import for subset: ${subset} completed.`);
+  } catch (error) {
+    console.error(`Error during dataset import for subset: ${subset}`, error);
+  }
+}
+
+// The API Route Handler
+export async function POST(request: Request) {
+  try {
+    const { subset } = await request.json();
+
+    if (!subset) {
+      return NextResponse.json({ error: 'Missing "subset" parameter (e.g., "python")' }, { status: 400 });
+    }
+
+    // Trigger the background job without waiting for it to complete
+    importDataset(subset).catch(console.error);
+
+    return NextResponse.json(
+      { message: `Started importing the "${subset}" subset from "${DATASET}". This will take a while.` },
+      { status: 202 }
+    );
+  } catch (error: any) {
+    console.error('API Error:', error);
+    return NextResponse.json({ error: 'Internal Server Error', detail: error.message }, { status: 500 });
+  }
+}

app/api/import-dataset/route.ts

@@ -0,0 +1,50 @@
+import { createClient } from '@supabase/supabase-js';
+import { pipeline } from '@xenova/transformers';
+
+// Configuration
+const EMBEDDING_MODEL = 'Xenova/all-MiniLM-L6-v2';
+
+// Initialize Supabase client
+const supabase = createClient(
+  process.env.NEXT_PUBLIC_SUPABASE_URL!,
+  process.env.SUPABASE_SERVICE_ROLE_KEY!
+);
+
+// Singleton pattern for the embedding pipeline
+let embeddingPipeline: any = null;
+async function getEmbeddingPipeline() {
+  if (!embeddingPipeline) {
+    embeddingPipeline = await pipeline('feature-extraction', EMBEDDING_MODEL);
+  }
+  return embeddingPipeline;
+}
+
+// Function to retrieve relevant code snippets
+export async function retrieveRelevantCode(query: string, match_threshold = 0.75, match_count = 5) {
+  try {
+    const pipeline = await getEmbeddingPipeline();
+
+    // 1. Generate an embedding for the user's query
+    const queryEmbedding = await pipeline(query, {
+      pooling: 'mean',
+      normalize: true,
+    });
+
+    // 2. Query Supabase for similar embeddings
+    const { data, error } = await supabase.rpc('match_code_embeddings', {
+      query_embedding: Array.from(queryEmbedding.data),
+      match_threshold,
+      match_count,
+    });
+
+    if (error) {
+      console.error('Supabase RPC error:', error);
+      return [];
+    }
+
+    return data.map((item: any) => item.content);
+  } catch (error) {
+    console.error('Error during code retrieval:', error);
+    return [];
+  }
+}

lib/retrieval.ts

@@ -1,4 +1,9 @@
 /** @type {import('next').NextConfig} */
-const nextConfig = {}
+const nextConfig = {
+  webpack: (config) => {
+    config.externals.push('onnxruntime-node')
+    return config
+  },
+}
 
 export default nextConfig