vinmay1234/SynapseEd

1 2 3from datasets import load_dataset4 5# Load 70% of the Wikipedia dataset6# dataset = load_dataset('wikimedia/wikipedia', "20231101.en", split='train[:70%]')7 8dataset = load_dataset('lucadiliello/wikipedia_512_pretraining',split = 'train[:70%]')9# from transformers import AutoModelForCausalLM, AutoTokenizer, BitsAndBytesConfig10 11# # Define the quantization configuration for 4-bit12# quantization_config = BitsAndBytesConfig(13#     load_in_4bit=True,              # Enable 4-bit precision14#     bnb_4bit_quant_type="nf4",      # Use the NF4 quantization type (good for reducing memory)15#     bnb_4bit_use_double_quant=True, # Enables double quantization to improve accuracy16#     bnb_4bit_compute_dtype="float16" # Use float16 for faster computation17# )18 19# # Load the tokenizer20# tokenizer = AutoTokenizer.from_pretrained('TinyLlama/TinyLlama-1.1B-Chat-v1.0')21 22# # Load the model with the quantization configuration23# model = AutoModelForCausalLM.from_pretrained(24#     'TinyLlama/TinyLlama-1.1B-Chat-v1.0',25#     quantization_config=quantization_config,  # Apply the 4-bit quantization config26#     device_map='auto'  # Automatically map model to available devices (e.g., GPU/CPU)27# )28 29# # Enable gradient checkpointing to reduce memory usage during training30# model.gradient_checkpointing_enable()31 32 33 34###########################################################    gpt2    ####################################################35 36 37 38from transformers import AutoModelForCausalLM, AutoTokenizer, BitsAndBytesConfig39 40# Define the quantization configuration for 4-bit41quantization_config = BitsAndBytesConfig(42    load_in_4bit=True,              # Enable 4-bit precision43    bnb_4bit_quant_type="nf4",      # Use the NF4 quantization type (good for reducing memory)44    bnb_4bit_use_double_quant=True, # Enables double quantization to improve accuracy45    bnb_4bit_compute_dtype="float16" # Use float16 for faster computation46)47 48# Load the tokenizer49tokenizer = AutoTokenizer.from_pretrained('gpt2')50 51# Load the model with the quantization configuration52model = AutoModelForCausalLM.from_pretrained(53    'gpt2',54    quantization_config=quantization_config,  # Apply the 4-bit quantization config55    device_map='auto'  # Automatically map model to available devices (e.g., GPU/CPU)56)57 58# Enable gradient checkpointing to reduce memory usage during training59model.gradient_checkpointing_enable()60 61 62 63from peft import LoraConfig, get_peft_model64import bitsandbytes as bnb65 66# Configure PEFT with 4-bit precision67# lora_config = LoraConfig(r=16, lora_alpha=32, target_modules=["q_proj", "v_proj"], lora_dropout=0.05, bias="none")68lora_config = LoraConfig(r=16, lora_alpha=32, target_modules=["attn.c_attn", "mlp.c_fc", "mlp.c_proj"], lora_dropout=0.05, bias="none")69peft_model = get_peft_model(model, lora_config)70 71 72# Set the pad token (using eos_token or adding a new special token)73if tokenizer.pad_token is None:74    # Option 1: Use eos_token as pad_token75    tokenizer.pad_token = tokenizer.eos_token76 77    # Option 2: Add [PAD] as a new pad token if needed78    # tokenizer.add_special_tokens({'pad_token': '[PAD]'})79 80# Tokenize the dataset with optimized settings81def tokenize_function(examples):82    return tokenizer(examples['text'], truncation=True, padding='max_length', max_length=150)83 84tokenized_dataset = dataset.select(range(100000)).map(tokenize_function, batched=True)85def prepare_labels(batch):86    batch["labels"] = batch["input_ids"].copy()  # Copy input_ids as labels for language modeling87    return batch88 89# Apply the transformation to add labels90tokenized_dataset = tokenized_dataset.map(prepare_labels, batched=True)91# Step 1: Install FAISS for the Vector Database92 93from datasets import Dataset94from transformers import AutoModel, AutoTokenizer95import faiss96import numpy as np97from tqdm import tqdm  # Import tqdm for progress bar98 99# Load your tokenizer and model100embedding_model_name = "sentence-transformers/all-mpnet-base-v2"101embedding_model = AutoModel.from_pretrained(embedding_model_name)102embedding_tokenizer = AutoTokenizer.from_pretrained(embedding_model_name)103 104# Move the model to GPU if available105device = torch.device("cuda" if torch.cuda.is_available() else "cpu")106embedding_model.to(device)107 108# Function to generate embeddings in batches109def embed_text_batch(texts, batch_size=16):110    all_embeddings = []111 112    for i in tqdm(range(0, len(texts), batch_size), desc="Generating embeddings"):113        batch_texts = texts[i:i + batch_size]114 115        # Tokenize and move inputs to the GPU116        inputs = embedding_tokenizer(batch_texts, padding=True, truncation=True, return_tensors="pt").to(device)117 118        with torch.no_grad():119            # Generate embeddings and move them back to CPU120            embeddings = embedding_model(**inputs).last_hidden_state.mean(dim=1).cpu().numpy()  # Mean pooling121 122        all_embeddings.extend(embeddings)123 124    return np.array(all_embeddings)125 126# Step 1: Process the dataset in batches127texts = tokenized_dataset["text"]128batch_size = 16  # Adjust based on Colab memory129embeddings = embed_text_batch(texts, batch_size=batch_size)130 131# Step 2: Add embeddings as a new column to the dataset132tokenized_dataset = tokenized_dataset.add_column("embeddings", embeddings.tolist())133 134# Step 3: Add FAISS index135dimension = embeddings.shape[1]  # Dimension of embeddings136faiss_index = faiss.IndexFlatL2(dimension)137 138# Step 4: Add embeddings to FAISS index139faiss_index.add(embeddings)140 141# Step 5: Save the dataset and FAISS index142tokenized_dataset.save_to_disk("wikipedia_dataset_with_embeddings")143faiss.write_index(faiss_index, "wikipedia_faiss.index")144 145print("FAISS index and dataset saved successfully.")146def embed_query(query):147    # Tokenize and embed the query148    inputs = embedding_tokenizer([query], padding=True, truncation=True, return_tensors="pt").to(device)149 150    with torch.no_grad():151        query_embedding = embedding_model(**inputs).last_hidden_state.mean(dim=1).cpu().numpy()152 153    return query_embedding154def search_faiss(query_embedding, faiss_index, top_k=5):155    # Search the FAISS index156    distances, indices = faiss_index.search(query_embedding, top_k)157 158    return distances, indices159def get_top_answer(indices, dataset):160    # Retrieve the top answer(s) from the dataset based on the indices161    return dataset["text"][indices[0][0]]  # Assuming top result, can adjust for more answers162import torch163from transformers import AutoTokenizer, AutoModelForSeq2SeqLM164import faiss165import numpy as np166 167# Assuming embeddings and faiss_index are already created as in your previous code168 169# Load the pre-trained LLM for generation (you can replace it with a different one)170llm_model_name = "facebook/bart-large-cnn"  # Example: You can use GPT-3, BART, T5, etc.171llm_model = AutoModelForSeq2SeqLM.from_pretrained(llm_model_name)172llm_tokenizer = AutoTokenizer.from_pretrained(llm_model_name)173 174# Move model to GPU if available175device = torch.device("cuda" if torch.cuda.is_available() else "cpu")176llm_model.to(device)177 178# Embedding model used for creating the vector database (same as the one used to generate embeddings for dataset)179embedding_model_name = "sentence-transformers/all-mpnet-base-v2"180embedding_tokenizer = AutoTokenizer.from_pretrained(embedding_model_name)181embedding_model = AutoModel.from_pretrained(embedding_model_name)182embedding_model.to(device)183 184# Function to embed a query (same as before)185def embed_query(query):186    inputs = embedding_tokenizer([query], padding=True, truncation=True, return_tensors="pt").to(device)187    with torch.no_grad():188        query_embedding = embedding_model(**inputs).last_hidden_state.mean(dim=1).cpu().numpy()189    return query_embedding190 191# Function to search FAISS index and retrieve top k results192def search_faiss(query_embedding, faiss_index, top_k=5):193    distances, indices = faiss_index.search(query_embedding, top_k)194    return distances, indices195 196# Function to generate an answer using the LLM based on the retrieved documents197def generate_answer(query, retrieved_texts):198    # Combine the query and the retrieved texts into a single input199    context = " ".join(retrieved_texts)200    input_text = f"Question: {query}\nContext: {context}\nAnswer:"201    202    # Tokenize and pass to the LLM203    inputs = llm_tokenizer(input_text, return_tensors="pt", max_length=512, truncation=True).to(device)204    with torch.no_grad():205        generated_ids = llm_model.generate(inputs['input_ids'], max_length=150)206    207    # Decode the generated response208    answer = llm_tokenizer.decode(generated_ids[0], skip_special_tokens=True)209    return answer210 211# Function to retrieve the texts from the dataset based on FAISS index results212def get_retrieved_texts(indices, dataset, top_k=5):213    retrieved_texts = []214    for idx in indices[0][:top_k]:  # Get the top K results215        retrieved_texts.append(dataset['text'][idx])  # Assuming 'text' is the relevant field in the dataset216    return retrieved_texts217 218# Example usage219def rag_pipeline(question, faiss_index, dataset, top_k=3):220    # Step 1: Embed the query221    query_embedding = embed_query(question)222 223    # Step 2: Search the FAISS index for the top K similar documents224    distances, indices = search_faiss(query_embedding, faiss_index, top_k=top_k)225 226    # Step 3: Retrieve the top K relevant documents from the dataset227    retrieved_texts = get_retrieved_texts(indices, dataset, top_k=top_k)228 229    # Step 4: Generate the answer using the retrieved texts and the LLM230    answer = generate_answer(question, retrieved_texts)231 232    return answer233 234# Import the necessary modules235from langchain_community.llms import Ollama236 237# Load the Ollama model238gen_model = Ollama(model="llama2")239 240# Define a function to get predefined responses for specific queries241def get_predefined_response(question):242    predefined_responses = {243        "hi": "Hello! How can I assist you today?",244        "hello": "Hi there! 😊 What can I help you with?",245        "who made you?": "I was created by Vinmay and his team.",246        "what is your purpose?": "I'm here to assist you with educational queries and provide information.",247        # Add more predefined responses as needed248    }249    250    # Normalize the question to make it case insensitive251    normalized_question = question.lower()252    253    return predefined_responses.get(normalized_question, None)254 255# Modify the generate_response function to check for predefined responses256def generate_response(markdown, question, user_instructions=None, max_new_tokens=250, temperature=0.9, top_p=0.95):257    # Check for predefined response first258    predefined_response = get_predefined_response(question)259    if predefined_response:260        return predefined_response261    262    instruction_text = f" Please follow these instructions: {user_instructions}" if user_instructions else ""263    264    prompt = (265        f"Using the provided context, please generate a unique and insightful answer that directly addresses the question:\n\n"266        f"Context:\n{markdown}\n\n"267        f"Question: {question}\n"268        f"{instruction_text}\n"269        f"If any personal query asked then refer{predefined_response}\n and based upon it, genarate your own answer"270        f"Please synthesize your response by integrating the information with your own understanding: "271    )272 273    # Call the Ollama model using the `invoke` method274    response = gen_model.invoke(prompt, max_tokens=max_new_tokens, temperature=temperature, top_p=top_p)275 276    # Check if the response is a string (direct generated text) or a dictionary (with metadata)277    if isinstance(response, str):278        return response  # Return the raw text if it's a string279    elif isinstance(response, dict) and "choices" in response:280        return response["choices"][0]["text"]  # Extract the text from the structured response281    else:282        return "Unexpected response format."283 284# # Example usage285# markdown = "The sky appears blue due to the scattering of light by the atmosphere."286# question = "Hi"287# response = generate_response(markdown, question)288 289# print(f"Model Response: {response}")290 291import gradio as gr292from langchain_community.llms import Ollama293 294# Load the Ollama model295gen_model = Ollama(model="llama2")296 297# Define the manual responses298manual_responses = {299    "hi": "Hello! How can I assist you today?",300    "hello": "Hi there! What would you like to know?",301    "who made you?": "I was created by OpenAI.",302    "what is your purpose?": "I'm here to assist with educational queries!"303}304 305# Function to generate responses306def generate_response(user_input):307    # Normalize user input for matching308    normalized_input = user_input.lower().strip()309 310    # Check for manual responses311    if normalized_input in manual_responses:312        return manual_responses[normalized_input]313 314    # For other questions, generate a response using the model315    prompt = f"Please provide a detailed answer to the following question:\n\nQuestion: {user_input}\n"316    317    response = gen_model.invoke(prompt)318    return response.strip()319 320# Create the Gradio interface321iface = gr.Interface(322    fn=generate_response,323    inputs=gr.Textbox(label="Ask a Question"),324    outputs=gr.Textbox(label="Response"),325    title="Q&A System",326    description="Ask me anything and I will respond accordingly."327)328 329# Launch the Gradio app330if __name__ == "__main__":331    iface.launch(share=True, inline = False)  # Use share=True to make it public if needed332