Spaces:
Sleeping
Sleeping
| import gradio as gr | |
| import faiss | |
| import numpy as np | |
| import openai | |
| from sentence_transformers import SentenceTransformer | |
| from nltk.tokenize import sent_tokenize | |
| # Load the Ubuntu manual from a .txt file | |
| with open("/content/ubuntu_manual.txt", "r", encoding="utf-8") as file: | |
| full_text = file.read() | |
| # Function to chunk the text into smaller pieces | |
| def chunk_text(text, chunk_size=500): # Larger chunks | |
| sentences = sent_tokenize(text) | |
| chunks = [] | |
| current_chunk = [] | |
| for sentence in sentences: | |
| if len(current_chunk) + len(sentence.split()) <= chunk_size: | |
| current_chunk.append(sentence) | |
| else: | |
| chunks.append(" ".join(current_chunk)) | |
| current_chunk = [sentence] | |
| if current_chunk: | |
| chunks.append(" ".join(current_chunk)) | |
| return chunks | |
| # Apply chunking to the entire text | |
| manual_chunks = chunk_text(full_text, chunk_size=500) | |
| # Load your FAISS index | |
| index = faiss.read_index("path/to/your/faiss_index.bin") | |
| # Load your embedding model | |
| embedding_model = SentenceTransformer('your_embedding_model_name') | |
| # OpenAI API key | |
| openai.api_key = 'your-openai-api-key' | |
| # Function to create embeddings | |
| def embed_text(text_list): | |
| return np.array(embedding_model.encode(text_list), dtype=np.float32) | |
| # Function to retrieve relevant chunks for a user query | |
| def retrieve_chunks(query, k=5): | |
| query_embedding = embed_text([query]) | |
| # Search the FAISS index | |
| distances, indices = index.search(query_embedding, k=k) | |
| # Debugging: Print out the distances and indices | |
| print("Distances:", distances) | |
| print("Indices:", indices) | |
| # Check if indices are valid | |
| if len(indices[0]) == 0: | |
| return [] | |
| # Ensure indices are within bounds | |
| valid_indices = [i for i in indices[0] if i < len(manual_chunks)] | |
| if not valid_indices: | |
| return [] | |
| # Retrieve relevant chunks | |
| relevant_chunks = [manual_chunks[i] for i in valid_indices] | |
| return relevant_chunks | |
| # Function to truncate long inputs | |
| def truncate_input(text, max_length=512): | |
| tokens = generator_tokenizer.encode(text, truncation=True, max_length=max_length, return_tensors="pt") | |
| return tokens | |
| # Function to perform RAG: Retrieve chunks and generate a response | |
| def rag_response(query, k=5, max_new_tokens=150): | |
| # Step 1: Retrieve relevant chunks | |
| relevant_chunks = retrieve_chunks(query, k=k) | |
| if not relevant_chunks: | |
| return "Sorry, I couldn't find relevant information." | |
| # Step 2: Combine the query with retrieved chunks | |
| augmented_input = query + "\n" + "\n".join(relevant_chunks) | |
| # Truncate and encode the input | |
| inputs = truncate_input(augmented_input) | |
| # Generate response | |
| outputs = generator_model.generate(inputs, max_new_tokens=max_new_tokens) | |
| generated_text = generator_tokenizer.decode(outputs[0], skip_special_tokens=True) | |
| return generated_text | |
| # Gradio Interface | |
| iface = gr.Interface( | |
| fn=rag_response, | |
| inputs="text", | |
| outputs="text", | |
| title="RAG Chatbot with FAISS and GPT-3.5", | |
| description="Ask me anything!" | |
| ) | |
| if __name__ == "__main__": | |
| iface.launch() | |