Initialize app

README.md

@@ -1,8 +1,8 @@
 ---
 title: Kubectl V2
 emoji: 🌍
-colorFrom: purple
-colorTo: purple
+colorFrom: blue
+colorTo: gray
 sdk: gradio
 sdk_version: 4.4.1
 app_file: app.py

app.py

@@ -0,0 +1,204 @@
+import gradio as gr
+from transformers.generation.stopping_criteria import StoppingCriteria, StoppingCriteriaList
+from transformers import AutoTokenizer, AutoModelForCausalLM, BitsAndBytesConfig
+from peft import PeftModel
+import torch
+import pinecone
+from sentence_transformers import SentenceTransformer
+from tqdm import tqdm
+from sentence_transformers.cross_encoder import CrossEncoder
+import numpy as np
+from torch import nn
+
+# Set up semantic search
+PINECONE_API_KEY = $PINECONE_API_KEY
+
+def get_embedding(text):
+  embed_text = sentencetransformer_model.encode(text)
+  vector_text = embed_text.tolist()
+
+  return vector_text
+
+def query_from_pinecone(query, top_k=3):
+    # get embedding from THE SAME embedder as the documents
+    query_embedding = get_embedding(query)
+
+    return index.query(
+      vector=query_embedding,
+      top_k=top_k,
+      include_metadata=True   # gets the metadata (dates, text, etc)
+    ).get('matches')
+
+def get_results_from_pinecone(query, top_k=3, re_rank=True, verbose=True):
+
+    results_from_pinecone = query_from_pinecone(query, top_k=top_k)
+    if not results_from_pinecone:
+        return []
+
+    if verbose:
+        print("Query:", query)
+
+
+    final_results = []
+
+    if re_rank:
+        if verbose:
+            print('Document ID (Hash)\t\tRetrieval Score\tCE Score\tText')
+
+        sentence_combinations = [[query, result_from_pinecone['metadata']['text']] for result_from_pinecone in results_from_pinecone]
+
+        # Compute the similarity scores for these combinations
+        similarity_scores = cross_encoder.predict(sentence_combinations, activation_fct=nn.Sigmoid())
+
+        # Sort the scores in decreasing order
+        sim_scores_argsort = reversed(np.argsort(similarity_scores))
+
+        # Print the scores
+        for idx in sim_scores_argsort:
+            result_from_pinecone = results_from_pinecone[idx]
+            final_results.append(result_from_pinecone)
+            if verbose:
+                print(f"{result_from_pinecone['id']}\t{result_from_pinecone['score']:.2f}\t{similarity_scores[idx]:.2f}\t{result_from_pinecone['metadata']['text'][:50]}")
+        return final_results
+
+    if verbose:
+        print('Document ID (Hash)\t\tRetrieval Score\tText')
+    for result_from_pinecone in results_from_pinecone:
+        final_results.append(result_from_pinecone)
+        if verbose:
+            print(f"{result_from_pinecone['id']}\t{result_from_pinecone['score']:.2f}\t{result_from_pinecone['metadata']['text'][:50]}")
+
+    return final_results
+
+def semantic_search(prompt):
+    final_results = get_results_from_pinecone(prompt, top_k=3, re_rank=True, verbose=True)
+
+    return 'First result:\n' + final_results[0]['metadata']['text'].replace('\n', ' ') + '\n' + 'Second result:\n' + final_results[1]['metadata']['text'].replace('\n', ' ') + '\n' + 'Third result:\n' + final_results[2]['metadata']['text'].replace('\n', ' ')
+
+cross_encoder = CrossEncoder('cross-encoder/ms-marco-MiniLM-L-12-v2')
+sentencetransformer_model = SentenceTransformer('sentence-transformers/multi-qa-mpnet-base-cos-v1')
+
+pinecone_key = PINECONE_API_KEY
+
+INDEX_NAME = 'k8s-semantic-search'
+NAMESPACE = 'default'
+
+pinecone.init(api_key=pinecone_key, environment="gcp-starter")
+
+if not INDEX_NAME in pinecone.list_indexes():
+    pinecone.create_index(
+        INDEX_NAME,  # The name of the index
+        dimension=768,  # The dimensionality of the vectors
+        metric='cosine',  # The similarity metric to use when searching the index
+        pod_type='starter'  # The type of Pinecone pod
+    )
+
+index = pinecone.Index(INDEX_NAME)
+
+# Set up mistral model
+base_model_id = 'mistralai/Mistral-7B-Instruct-v0.1'
+lora_model_id = 'ComponentSoft/mistral-kubectl-instruct'
+
+tokenizer = AutoTokenizer.from_pretrained(
+    lora_model_id,
+    padding_side="left",
+    add_eos_token=False,
+    add_bos_token=True,
+)
+tokenizer.pad_token = tokenizer.eos_token
+
+bnb_config = BitsAndBytesConfig(
+    load_in_4bit=True,
+    bnb_4bit_use_double_quant=True,
+    bnb_4bit_quant_type="nf4",
+    bnb_4bit_compute_dtype=torch.bfloat16
+)
+
+base_model = AutoModelForCausalLM.from_pretrained(
+    base_model_id,
+    quantization_config=bnb_config,
+    use_cache=True,
+    trust_remote_code=True,
+)
+
+model = PeftModel.from_pretrained(base_model, lora_model_id)
+model.eval()
+
+stop_terms=["</s>", "#End"]
+eos_token_ids_custom =  [torch.tensor(tokenizer.encode(term, add_special_tokens=False)).to("cuda") for term in stop_terms]
+
+category_terms=["</s>", "\n"]
+category_eos_token_ids_custom =  [torch.tensor(tokenizer.encode(term, add_special_tokens=False)).to("cuda") for term in stop_terms]
+
+
+class EvalStopCriterion(StoppingCriteria):
+    def __call__(self, input_ids: torch.LongTensor, score: torch.FloatTensor, **kwargs):
+        return any(torch.equal(e, input_ids[0][-len(e):]) for e in eos_token_ids_custom)
+
+
+class CategoryStopCriterion(StoppingCriteria):
+    def __call__(self, input_ids: torch.LongTensor, score: torch.FloatTensor, **kwargs):
+        return any(torch.equal(e, input_ids[0][-len(e):]) for e in category_eos_token_ids_custom)
+
+start_template = '### Answer:'
+command_template = '# Command:'
+end_template = '#End'
+
+def text_to_text_generation(prompt):
+    prompt = prompt.strip()
+    ''
+
+    is_kubectl_prompt = (
+        f"[INST] You are a helpful assistant who classifies prompts into three categories. Respond with 0 if it pertains to a 'kubectl' operation. This is an instruction that can be answered with a 'kubectl' action. Look for keywords like 'get', 'list', 'create', 'show', 'view', and other command-like words. This category is an instruction instead of a question. Respond with 1 only if the prompt is a question, and is about a definition related to Kubernetes, or non-action inquiries. Respond with 2 every other scenario, for example if the question is a general question, not related to Kubernetes or 'kubectl'.\n"
+        f"So for instance the following:\n"
+        f"List all pods in Kubernetes\n"
+        f"Would get a response:\n"
+        f"0 [/INST]"
+        f'text: "{prompt}"'
+        f'response (0/1/2): '
+    )
+
+
+    model_input = tokenizer(is_kubectl_prompt, return_tensors="pt").to("cuda")
+    with torch.no_grad():
+        response = tokenizer.decode(model.generate(**model_input, max_new_tokens=8, pad_token_id=tokenizer.eos_token_id, repetition_penalty=1.15, stopping_criteria=StoppingCriteriaList([CategoryStopCriterion()]))[0], skip_special_tokens=True)
+    response = response[len(is_kubectl_prompt):]
+
+    print('-----------------------------QUERY START-----------------------------')
+    print('Prompt: ' + prompt)
+    print('Classified as: ' + response)
+    response_num = 2 # Default to generic question
+    if '0' in response:
+      response_num = 0
+    elif '1' in response:
+      response_num = 1
+
+
+    # Check if general question
+    if response_num == 0:
+        prompt = f'[INST] {prompt}\n Lets think step by step. [/INST] {start_template}'
+    elif response_num == 1:
+        retrieved_results = semantic_search(prompt)
+        print('Query:')
+        print(f'[INST] You are an assistant who summarizes results retrieved from a book about Kubernetes. This summary should answer the question. If the answer is not in the retrieved results, use your general knowledge. [/INST] Question: {prompt}\nRetrieved results:\n{retrieved_results}\nResponse:')
+        prompt = f'[INST] You are an assistant who summarizes results retrieved from a book about Kubernetes. This summary should answer the question. If the answer is not in the retrieved results, use your general knowledge. [/INST] Question: {prompt}\nRetrieved results:\n{retrieved_results}\nResponse:'
+    else:
+        prompt = f'[INST] {prompt} [/INST]'
+
+    # Generate output
+    model_input = tokenizer(prompt, return_tensors="pt").to("cuda")
+    with torch.no_grad():
+      response = tokenizer.decode(model.generate(**model_input, max_new_tokens=256, pad_token_id=tokenizer.eos_token_id, repetition_penalty=1.15, stopping_criteria=StoppingCriteriaList([EvalStopCriterion()]))[0], skip_special_tokens=True)
+
+    # Get the relevalt parts
+    start = response.index(start_template) + len(start_template) if start_template in response else len(prompt)
+    start = response.index(command_template) + len(command_template) if command_template in response else start
+    end = response.index(end_template) if end_template in response else len(response)
+    true_response = response[start:end].strip()
+    print('Returned: ' + true_response)
+    print('------------------------------QUERY END------------------------------')
+
+    return true_response
+
+iface = gr.Interface(fn=semantic_search, inputs="text", outputs="text")
+iface.launch()

requirements.txt

@@ -0,0 +1,8 @@
+transformers
+peft
+bitsandbytes
+torch
+scipy
+pinecone-client
+sentence_transformers
+tqdm