Upload 3 files

app.py

@@ -1,147 +1,140 @@
-import os
-import chromadb
-import gradio as gr
-from transformers import AutoModelForSeq2SeqLM, AutoTokenizer, pipeline
-from langchain_chroma import Chroma
-from langchain_huggingface import HuggingFaceEmbeddings, HuggingFacePipeline
-from langchain_community.document_loaders import PyPDFLoader
-from langchain.text_splitter import RecursiveCharacterTextSplitter
-from langchain.chains import create_retrieval_chain, LLMChain
-from langchain.prompts import PromptTemplate
-from collections import OrderedDict
-
-# Load embeddings model
-embeddings = HuggingFaceEmbeddings(model_name="sentence-transformers/all-MiniLM-L6-v2")
-
-# Load Chroma database (Avoid reprocessing documents)
-CHROMA_PATH = "./chroma_db"
-if not os.path.exists(CHROMA_PATH):
-    raise FileNotFoundError("ChromaDB folder not found. Make sure it's uploaded to the repo.")
-
-chroma_client = chromadb.PersistentClient(path=CHROMA_PATH)
-db = Chroma(embedding_function=embeddings, client=chroma_client)
-
-# Load the model
-model_name = "google/flan-t5-large"
-tokenizer = AutoTokenizer.from_pretrained(model_name)
-model = AutoModelForSeq2SeqLM.from_pretrained(model_name)
-
-# Create pipeline
-qa_pipeline = pipeline(
-    "text2text-generation",
-    model=model,
-    tokenizer=tokenizer,
-    device=0,
-    max_length=512,
-    min_length=50,
-    do_sample=False,
-    repetition_penalty=1.2
-)
-
-# Wrap pipeline in LangChain
-llm = HuggingFacePipeline(pipeline=qa_pipeline)
-retriever = db.as_retriever(search_kwargs={"k": 3})
-
-
-def clean_context(context_list, max_tokens=350, min_length=50):
-    """
-    Cleans retrieved document context:
-    - Removes duplicates while preserving order
-    - Limits total token count
-    - Ensures useful, non-repetitive context
-    """
-
-    # Preserve order while removing duplicates
-    unique_texts = list(OrderedDict.fromkeys([doc.page_content.strip() for doc in context_list]))
-
-    # Remove very short texts (e.g., headers)
-    filtered_texts = [text for text in unique_texts if len(text.split()) > min_length]
-
-    # Avoid near-duplicate entries
-    deduplicated_texts = []
-    seen_texts = set()
-    for text in filtered_texts:
-        if not any(text in s for s in seen_texts):  # Avoid near-duplicates
-            deduplicated_texts.append(text)
-            seen_texts.add(text)
-
-    # Limit context based on token count
-    trimmed_context = []
-    total_tokens = 0
-    for text in deduplicated_texts:
-        tokenized_text = tokenizer.encode(text, add_special_tokens=False)
-        token_count = len(tokenized_text)
-
-        if total_tokens + token_count > max_tokens:
-            remaining_tokens = max_tokens - total_tokens
-            if remaining_tokens > 20:
-                trimmed_context.append(tokenizer.decode(tokenized_text[:remaining_tokens]))
-            break
-
-        trimmed_context.append(text)
-        total_tokens += token_count
-
-    return "\n\n".join(trimmed_context) if trimmed_context else "No relevant context found."
-
-# Define prompt
-prompt_template = PromptTemplate(
-    template="""
-    You are a Kubernetes instructor. Answer the question based on the provided context.
-    If the context does not provide an answer, say "I don't have enough information."
-
-    Context:
-    {context}
-
-    Question:
-    {input}
-
-    Answer:
-    """,
-    input_variables=["context", "input"]
-)
-
-llm_chain = LLMChain(llm=llm, prompt=prompt_template)
-qa_chain = create_retrieval_chain(retriever, llm_chain)
-
-# Query function
-def get_k8s_answer(query):
-    retrieved_context = retriever.get_relevant_documents(query)
-    cleaned_context = clean_context(retrieved_context, max_tokens=350)  # Ensure context size is within limits
-
-    # Ensure total input tokens < 512 before passing to model
-    input_text = f"Context:\n{cleaned_context}\n\nQuestion: {query}\nAnswer:"
-    total_tokens = len(tokenizer.encode(input_text, add_special_tokens=True))
-
-    if total_tokens > 512:
-      # Trim context further to fit within the limit
-      allowed_tokens = 512 - len(tokenizer.encode(query, add_special_tokens=True)) - 50  # 50 tokens for the model's response
-      cleaned_context = clean_context(retrieved_context, max_tokens=allowed_tokens)
-
-      # Recalculate total tokens
-      input_text = f"Context:\n{cleaned_context}\n\nQuestion: {query}\nAnswer:"
-      total_tokens = len(tokenizer.encode(input_text, add_special_tokens=True))
-
-      if total_tokens > 512:
-          return "Error: Even after trimming, input is too large."
-
-    response = qa_chain.invoke({"input": query, "context": cleaned_context})
-    return response
-
-def get_k8s_answer_text(query):
-  model_full_answer = get_k8s_answer(query)
-  if 'answer' in model_full_answer.keys():
-    if 'text' in model_full_answer['answer'].keys():
-      return model_full_answer['answer']['text']
-  return "Error"
-
-# Gradio Interface
-demo = gr.Interface(
-    fn=get_k8s_answer_text,
-    inputs=gr.Textbox(label="Ask a Kubernetes Question"),
-    outputs=gr.Textbox(label="Answer"),
-    title="Kubernetes RAG Assistant",
-    description="Ask any Kubernetes-related question and get a step-by-step answer based on documentation."
-)
-
-if __name__ == "__main__":
-    demo.launch()
+import gradio as gr
+import torch
+from transformers import AutoModelForSeq2SeqLM, AutoTokenizer, pipeline
+import chromadb
+from langchain.vectorstores import Chroma
+from langchain_huggingface import HuggingFaceEmbeddings, HuggingFacePipeline
+from langchain.chains import create_retrieval_chain, LLMChain
+from langchain.prompts import PromptTemplate
+import os
+import shutil
+import zipfile
+
+# 🚀 Step 1: Extract ChromaDB if not already done (only once)
+if not os.path.exists("./chroma_db"):
+    with zipfile.ZipFile("chroma.zip", "r") as zip_ref:
+        zip_ref.extractall("./chroma_db")
+
+# 🚀 Step 2: Load Pre-trained Model & Tokenizer (Fast Startup)
+MODEL_NAME = "google/flan-t5-xl"
+tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)
+model = AutoModelForSeq2SeqLM.from_pretrained(MODEL_NAME)
+
+# 🚀 Step 3: Load Vector Store Efficiently
+embeddings = HuggingFaceEmbeddings(model_name="sentence-transformers/all-MiniLM-L6-v2")
+chroma_client = chromadb.PersistentClient(path="./chroma_db")
+db = Chroma(embedding_function=embeddings, client=chroma_client)
+
+# 🚀 Step 4: Optimize Retriever (Lower `k` for Speed)
+retriever = db.as_retriever(search_kwargs={"k": 10})
+
+# 🚀 Step 5: Define Prompt for the LLM
+prompt_template = PromptTemplate(
+    template="""
+        You are a Kubernetes expert.  
+    **Answer the question using ONLY the provided context.**  
+    If the context does NOT contain enough information, return:  
+    `"I don't have enough information to answer this question."`  
+    Always include YAML examples when relevant. 
+    
+    ---
+    **Context:**  
+    {context}  
+    
+    **Question:**  
+    {input}  
+    
+    ---
+    **Answer:**  
+    """,
+    input_variables=["context", "input"]
+)
+
+# 🚀 Step 6: Build Retrieval Chain
+qa_pipeline = pipeline("text2text-generation", model=model, tokenizer=tokenizer, device=0,
+                       max_length=512, min_length=50, do_sample=True, temperature=0.4, top_p=0.9)
+llm = HuggingFacePipeline(pipeline=qa_pipeline)
+llm_chain = LLMChain(llm=llm, prompt=prompt_template)
+qa_chain = create_retrieval_chain(retriever, llm_chain)
+
+# 🚀 Step 7: Define Fast Answer Function
+def clean_context(context_list, max_tokens=350, min_length=50):
+    """
+    Improves the retrieved document context:
+    - Removes duplicates while preserving order
+    - Filters out very short or unstructured text
+    - Limits token count for better LLM performance
+    """
+    from collections import OrderedDict
+
+    # Preserve order while removing exact duplicates
+    unique_texts = list(OrderedDict.fromkeys(doc.page_content.strip() for doc in context_list))
+
+    # Remove very short texts (e.g., headers, page numbers)
+    filtered_texts = [text for text in unique_texts if len(text.split()) > min_length]
+
+    # Avoid near-duplicates
+    deduplicated_texts = []
+    seen_texts = set()
+    for text in filtered_texts:
+        normalized_text = " ".join(text.split())  # Normalize spacing
+        if not any(normalized_text in seen for seen in seen_texts):  # Avoid near-duplicates
+            deduplicated_texts.append(normalized_text)
+            seen_texts.add(normalized_text)
+
+    # Limit context based on token count
+    trimmed_context = []
+    total_tokens = 0
+    for text in deduplicated_texts:
+        tokenized_text = tokenizer.encode(text, add_special_tokens=False)
+        token_count = len(tokenized_text)
+
+        if total_tokens + token_count > max_tokens:
+            remaining_tokens = max_tokens - total_tokens
+            if remaining_tokens > 20:  # Allow partial inclusion if meaningful
+                trimmed_context.append(tokenizer.decode(tokenized_text[:remaining_tokens]))
+            break
+
+        trimmed_context.append(text)
+        total_tokens += token_count
+
+    return "\n\n".join(trimmed_context) if trimmed_context else "No relevant context found."
+
+
+def get_k8s_answer(query):
+    retrieved_context = retriever.invoke(query)
+    cleaned_context = clean_context(retrieved_context, max_tokens=350)
+
+    # Ensure input tokens fit within 512 limit
+    input_text = f"Context:\n{cleaned_context}\n\nQuestion: {query}\nAnswer:"
+    total_tokens = len(tokenizer.encode(input_text, add_special_tokens=True))
+
+    if total_tokens > 512:
+        # Further trim context
+        allowed_tokens = 512 - len(tokenizer.encode(query, add_special_tokens=True)) - 50  # 50 tokens reserved for response
+        cleaned_context = clean_context(retrieved_context, max_tokens=allowed_tokens)
+
+        # Recalculate total tokens
+        input_text = f"Context:\n{cleaned_context}\n\nQuestion: {query}\nAnswer:"
+        total_tokens = len(tokenizer.encode(input_text, add_special_tokens=True))
+
+        if total_tokens > 512:
+            return "Error: Even after trimming, input is too large."
+
+    response = qa_chain.invoke({"input": query, "context": cleaned_context})
+    return response
+
+# 🚀 Step 8: Optimize Gradio App with `Blocks()`
+with gr.Blocks(theme="soft") as demo:
+    gr.Markdown("# ⚡ Kubernetes RAG")
+    gr.Markdown("Ask any Kubernetes-related question!")
+
+    with gr.Row():
+        question = gr.Textbox(label="Ask a Kubernetes Question", lines=1)
+        answer = gr.Textbox(label="Answer", interactive=False)
+
+    submit_button = gr.Button("Get Answer")
+
+    submit_button.click(fn=get_k8s_answer, inputs=question, outputs=answer)
+
+demo.launch()

chroma.zip

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:1cb847a7f5e922fead2197320f50734db00b6280e5acc8c202317b67f484e46a
+size 126566892

requirements.txt

@@ -1,11 +1,9 @@
-huggingface_hub==0.25.2
-gradio
-transformers
-sentence-transformers
-chromadb
-pypdf
-torch
-langchain
-langchain-huggingface
-langchain-chroma
-langchain_community
+torch
+transformers
+gradio
+chromadb
+langchain
+langchain-chroma
+langchain-community
+langchain-huggingface
+sentence-transformers