Upload 2 files

app.py

@@ -0,0 +1,151 @@
+import gradio as gr
+import torch
+import os
+from transformers import AutoModelForCausalLM, AutoTokenizer, pipeline
+from langchain_community.document_loaders import DirectoryLoader
+from langchain_text_splitters import RecursiveCharacterTextSplitter, Language
+from langchain_huggingface import HuggingFaceEmbeddings, HuggingFacePipeline
+from langchain_community.vectorstores import FAISS
+from langchain_core.runnables import RunnablePassthrough
+from langchain_core.output_parsers import StrOutputParser
+from langchain_core.prompts import PromptTemplate
+
+# 1. HARDWARE OPTIMIZED LLM LOADING
+
+def load_llm():
+    model_id = "Qwen/Qwen2.5-Coder-0.5B-Instruct"
+    tokenizer = AutoTokenizer.from_pretrained(model_id)
+    
+    model = AutoModelForCausalLM.from_pretrained(
+        model_id,
+        device_map="auto",
+        torch_dtype="auto",
+        low_cpu_mem_usage=True
+    )
+    
+    pipe = pipeline(
+        "text-generation",
+        model=model,
+        tokenizer=tokenizer,
+        max_new_tokens=300,
+        temperature=0.1,
+        repetition_penalty=1.1,
+        return_full_text=False
+    )
+    return HuggingFacePipeline(pipeline=pipe)
+
+# 2. CODE INGESTION & VECTOR DATABASE
+
+def setup_vector_db():
+    if not os.path.exists('./repo'):
+        os.makedirs('./repo')
+        
+    loader = DirectoryLoader('./repo', glob="**/*.py", show_progress=True)
+    docs = loader.load()
+    
+    if not docs:
+        return None, 0
+
+    python_splitter = RecursiveCharacterTextSplitter.from_language(
+        language=Language.PYTHON, 
+        chunk_size=500, 
+        chunk_overlap=50
+    )
+    texts = python_splitter.split_documents(docs)
+    
+    embeddings = HuggingFaceEmbeddings(model_name="all-MiniLM-L6-v2")
+    db = FAISS.from_documents(texts, embeddings)
+    
+    return db, len(docs)
+
+# 3. GLOBAL INITIALIZATION
+print("Initializing models...")
+device_status = "🟢 GPU Active" if torch.cuda.is_available() else "🟡 CPU Mode"
+llm = load_llm()
+vector_db, file_count = setup_vector_db()
+
+prompt_template = """Use the following codebase context to answer the question. 
+If you don't know the answer, just say that you don't know, don't try to make up code.
+
+Context: {context}
+
+Question: {input}
+Helpful Developer Answer:"""
+
+prompt = PromptTemplate.from_template(prompt_template)
+
+def format_docs(docs):
+    return "\n\n".join(doc.page_content for doc in docs)
+
+if vector_db:
+    retriever = vector_db.as_retriever(search_kwargs={"k": 3})
+    qa_chain = (
+        {"context": retriever, "input": RunnablePassthrough()}
+        | RunnablePassthrough.assign(
+            answer=(
+                RunnablePassthrough.assign(context=lambda x: format_docs(x["context"]))
+                | prompt
+                | llm
+                | StrOutputParser()
+            )
+        )
+    )
+else:
+    qa_chain = None
+
+# 4. CHAT LOGIC
+def respond(message, chat_history):
+    if not vector_db:
+        bot_message = "👋 Welcome! Please upload some Python files to the `./repo` directory and restart the server to start chatting."
+        chat_history.append((message, bot_message))
+        return "", chat_history
+    
+    # Fetch response from RAG
+    response = qa_chain.invoke(message)
+    answer = response["answer"]
+    sources = response["context"]
+    
+    final_answer = answer
+    
+    if sources:
+        final_answer += "\n\n<details><summary>🔍 View Source Code Referenced</summary>\n\n"
+        for idx, doc in enumerate(sources):
+            source_file = doc.metadata.get("source", "Unknown File")
+            final_answer += f"**Snippet {idx + 1}** from `{source_file}`:\n"
+            final_answer += f"```python\n{doc.page_content}\n```\n\n"
+        final_answer += "</details>"
+        
+    chat_history.append((message, final_answer))
+    return "", chat_history
+
+# 5. GRADIO UI
+custom_css = """
+.status-box { padding: 10px; border-radius: 8px; background-color: #f0f0f0; margin-bottom: 10px; }
+.dark .status-box { background-color: #1e293b; color: #cbd5e1; }
+"""
+
+with gr.Blocks(title="Codebase Assistant", css=custom_css) as demo:
+    with gr.Row():
+        with gr.Column(scale=1):
+            gr.Markdown("# DevAssist AI\nYour personal Qwen-powered codebase expert.")
+            gr.Markdown("---")
+            
+            with gr.Column(elem_classes=["status-box"]):
+                gr.Markdown("### System Status")
+                gr.Markdown(f"**Hardware:** {device_status}")
+                if vector_db:
+                    gr.Markdown(f"**Repo Status:** {file_count} files indexed ✅")
+                else:
+                    gr.Markdown("**Repo Status:** Empty ❌\n\nDrop your `.py` files into the `/repo` folder to begin analyzing.")
+        
+        with gr.Column(scale=3):
+            gr.Markdown("### 💻 Chat with your Codebase\nAsk architecture questions, find bugs, or request code explanations.")
+            chatbot = gr.Chatbot(height=500, show_label=False)
+            msg = gr.Textbox(placeholder="E.g., What does the main function do?", show_label=False)
+            clear = gr.Button("Clear Chat")
+
+            msg.submit(respond, inputs=[msg, chatbot], outputs=[msg, chatbot])
+            clear.click(lambda: None, None, chatbot, queue=False)
+
+if __name__ == "__main__":
+    demo.launch(server_name="0.0.0.0")

requirements.txt

@@ -0,0 +1,9 @@
+gradio
+torch
+transformers
+accelerate
+langchain
+langchain-community
+langchain-huggingface
+faiss-cpu
+sentence-transformers