Spaces:
Sleeping
Sleeping
| import faiss | |
| import numpy as np | |
| import os | |
| import streamlit as st | |
| from sentence_transformers import SentenceTransformer | |
| from transformers import AutoModelForCausalLM, AutoTokenizer | |
| # Initialize the Sentence-Transformer model for document embeddings | |
| embedding_model = SentenceTransformer('all-MiniLM-L6-v2') | |
| # Example Knowledge Base (can be expanded or replaced with real documents) | |
| #documents = [ | |
| # "The Industrial Revolution started in Britain in the 18th century.", | |
| # "Factories and machines revolutionized manufacturing processes.", | |
| # "Steam engines improved transportation and industrial production.", | |
| #"Capitalism and technological innovations shaped modern economies." | |
| #] | |
| # Path to the folder containing your content files | |
| content_file = "Ilango Profile.txt" | |
| # Read all text files in the folder | |
| # Read the document.txt file | |
| with open(content_file, "r") as file: | |
| documents = file.readlines() # This will load each line in the text file as a separate document | |
| # Convert documents to embeddings | |
| document_embeddings = embedding_model.encode(documents) | |
| # Set up FAISS for fast document retrieval | |
| dimension = document_embeddings.shape[1] | |
| index = faiss.IndexFlatL2(dimension) | |
| index.add(np.array(document_embeddings)) | |
| # Load GPT-2 model and tokenizer from Hugging Face | |
| model_name = "gpt2" | |
| tokenizer = AutoTokenizer.from_pretrained(model_name) | |
| model = AutoModelForCausalLM.from_pretrained(model_name) | |
| # Function to retrieve top-k relevant documents | |
| def retrieve_relevant_docs(query, top_k=3): | |
| query_embedding = embedding_model.encode([query]) # Get the embedding of the query | |
| _, idx = index.search(np.array(query_embedding), k=top_k) # Search for the top-k most relevant docs | |
| return "\n\n".join([documents[i] for i in idx[0]]) | |
| # Function to generate response using GPT-2 | |
| def generate_response(context, query): | |
| prompt = f"Context: {context}\n\nUser Query: {query}\n\nAnswer:" | |
| inputs = tokenizer(prompt, return_tensors="pt") | |
| outputs = model.generate(inputs['input_ids'], max_length=500, num_return_sequences=1) | |
| return tokenizer.decode(outputs[0], skip_special_tokens=True) | |
| # Streamlit UI | |
| st.title("Retrieval-Augmented Generation (RAG) System") | |
| st.write("Ask a question, and the system will retrieve relevant documents and generate an answer.") | |
| user_query = st.text_input("Enter your question:") | |
| if user_query: | |
| # Step 1: Retrieve relevant documents | |
| retrieved_docs = retrieve_relevant_docs(user_query) | |
| # Show retrieved documents | |
| st.write("#### Retrieved Documents:") | |
| st.write(retrieved_docs) | |
| # Step 2: Generate response based on retrieved documents | |
| response = generate_response(retrieved_docs, user_query) | |
| # Show the generated answer | |
| st.write("#### Answer:") | |
| st.write(response) |