| import pandas as pd |
| from .embedding import Embedder |
| from .retrieval import RetrievalDB |
| from .codegen import generate_code_with_context |
| class CodeGenPipeline: |
| def __init__(self, parquet_path): |
| self.df = pd.read_parquet(parquet_path) |
| self.prompts = self.df["prompt"].tolist() |
| self.solutions = self.df["canonical_solution"].tolist() |
| self.embedder = Embedder() |
| self.embeddings = self.embedder.encode(self.prompts, batch_size=32, show_progress_bar=True) |
| self.retrieval_db = RetrievalDB(self.prompts, self.embeddings, self.solutions) |
|
|
| def generate_code_from_prompt(self, user_prompt, k=1): |
| query_emb = self.embedder.encode([user_prompt])[0] |
| retrieved = self.retrieval_db.retrieve_similar_context(query_emb, k=k)[0] |
| context = "\n\n".join([f"# Task:\n{r['prompt']}\n{r['solution']}" for r in retrieved]) |
| return generate_code_with_context(user_prompt, context) |
