Clean repo without binary files

app.py

@@ -0,0 +1,39 @@
+import streamlit as st
+from pipeline.pipeline import CodeGenPipeline
+
+st.set_page_config(page_title="CodeGenBot", page_icon="🤖")
+st.title("💻 CodeGenBot")
+
+# Initialize pipeline only once (cache in session state)
+if "pipeline" not in st.session_state:
+    st.session_state.pipeline = CodeGenPipeline("hf://datasets/openai/openai_humaneval/openai_humaneval/test-00000-of-00001.parquet")
+
+# Memory for chat messages
+if "messages" not in st.session_state:
+    st.session_state.messages = []
+
+# Display previous messages
+for msg in st.session_state.messages:
+    if msg["role"] == "assistant":
+        st.chat_message("assistant").code(msg["content"], language="python")
+    else:
+        st.chat_message(msg["role"]).write(msg["content"])
+
+# User input
+user_input = st.chat_input("Ask CodeGenBot to generate Python code...")
+
+if user_input:
+    # Save user message
+    st.session_state.messages.append({"role": "user", "content": user_input})
+    st.chat_message("user").write(user_input)
+
+    # Generate code using your pipeline
+    with st.spinner("Generating code..."):
+        try:
+            code_output = st.session_state.pipeline.generate_code_from_prompt(user_input)
+        except Exception as e:
+            code_output = f"Error: {e}"
+
+    # Save assistant's reply and rerun to display it immediately
+    st.session_state.messages.append({"role": "assistant", "content": code_output})
+    st.rerun()

pipeline/codegen.py

@@ -0,0 +1,20 @@
+from dotenv import load_dotenv
+import os
+from huggingface_hub import InferenceClient
+
+load_dotenv()
+API_KEY = os.getenv("HF_API_KEY")
+MODEL_NAME = "deepseek-ai/DeepSeek-R1-Distill-Qwen-1.5B"
+
+
+def generate_code_with_context(user_prompt, context=None):
+    client = InferenceClient(provider="nscale", api_key=API_KEY)
+    if context:
+        final_prompt = f"{context}\n\n# Your Task:\n{user_prompt}\n"
+    else:
+        final_prompt = user_prompt
+    completion = client.chat.completions.create(
+        model=MODEL_NAME,
+        messages=[{"role": "user", "content": final_prompt}],
+    )
+    return completion.choices[0].message.content 

pipeline/embedding.py

@@ -0,0 +1,8 @@
+from sentence_transformers import SentenceTransformer
+
+class Embedder:
+    def __init__(self, model_name='all-MiniLM-L6-v2'):
+        self.model = SentenceTransformer(model_name)
+
+    def encode(self, texts, batch_size=32, show_progress_bar=False):
+        return self.model.encode(texts, batch_size=batch_size, show_progress_bar=show_progress_bar) 

pipeline/pipeline.py

@@ -0,0 +1,18 @@
+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) 

pipeline/retrieval.py

@@ -0,0 +1,16 @@
+import chromadb
+
+class RetrievalDB:
+    def __init__(self, prompts, embeddings, solutions, collection_name="humaneval"):
+        self.client = chromadb.Client()
+        self.collection = self.client.create_collection(name=collection_name)
+        for idx, (emb, prompt, solution) in enumerate(zip(embeddings, prompts, solutions)):
+            self.collection.add(
+                ids=[str(idx)],
+                embeddings=[emb.tolist()],
+                metadatas=[{"prompt": prompt, "solution": solution}]
+            )
+
+    def retrieve_similar_context(self, query_emb, k=1):
+        results = self.collection.query(query_embeddings=[query_emb], n_results=k)
+        return results["metadatas"] 

requirements.txt

@@ -0,0 +1,5 @@
+pandas
+sentence-transformers
+chromadb
+huggingface_hub
+pyarrow