Spaces:
Paused
Paused
| 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) | |