gitlabChatbot / app.py
jahnaviKolli's picture
Upload app.py
d6bff6a verified
import gradio as gr
import pickle
import faiss
import numpy as np
from sentence_transformers import SentenceTransformer
from transformers import pipeline
# Step 1: Load saved chunks
with open("chunks.pkl", "rb") as f:
chunks = pickle.load(f)
# Step 2: Load FAISS index
index = faiss.read_index("gitlab_index.faiss")
#Loading the embedding model
embedding_model = SentenceTransformer('all-MiniLM-L6-v2')
# Loading a Question-Answer model from Transformers pipeline
qa_pipeline = pipeline("question-answering", model="distilbert-base-uncased-distilled-squad")
# Define retrieval function
def query_knowledge_base(query, top_k=3):
query_embedding = embedding_model.encode([query])
distances, indices = index.search(np.array(query_embedding), top_k)
results = [chunks[i] for i in indices[0]]
return results
def generate_answer(context, question):
input_text = context if isinstance(context, str) else " ".join(context)
result = qa_pipeline(question=question, context=input_text)
return result["answer"]
def rag_chatbot(question):
try:
# Step 1: Embed the query and search the index
top_chunks = query_knowledge_base(question)
# Step 2: Combine top chunks into a single context string
context = " ".join(top_chunks)
result = generate_answer(context, question)
return result
except Exception as e:
return f"An error occurred: {str(e)}"
# Gradio Interface
def chat_interface_fn(message, history):
response = rag_chatbot(message)
return response # <-- Just return the response string
gr.ChatInterface(
fn=chat_interface_fn,
chatbot=gr.Chatbot(type='messages'),
title="GitLab Handbook Assistant",
description="Ask me about GitLab hiring or company policies!"
).launch(share=True, debug=True)