Update pages/4-Create your own bot (advanced).py

pages/4-Create your own bot (advanced).py

@@ -86,60 +86,92 @@ def developer_guide():
 
 # Call the function to display the developer guide page
 #developer_guide()
-
 import streamlit as st
-import PyPDF2
-from transformers import pipeline
+from transformers import RagRetriever, RagSequenceForGeneration, AutoTokenizer, AutoModelForSeq2SeqLM
+from PyPDF2 import PdfReader
 import torch
 
-# Hugging Face model
-MODEL_NAME = "facebook/rag-sequence-nq"
-
+# Load the tokenizer and the custom model (GPT-Neo-125M)
 @st.cache_resource
-def load_model():
-    # Using Hugging Face's pipeline for question-answering with DistilBERT
-    qa_pipeline = pipeline("question-answering", model="distilbert-base-uncased-distilled-squad")
-    return qa_pipeline
+def load_gpt_neo_rag():
+    tokenizer = AutoTokenizer.from_pretrained("EleutherAI/gpt-neo-125M")
+    
+    # Load GPT-Neo as the generator
+    custom_generator = AutoModelForSeq2SeqLM.from_pretrained("EleutherAI/gpt-neo-125M")
+    
+    # Initialize RAG retriever
+    retriever = RagRetriever.from_pretrained("facebook/rag-token-nq", index_name="exact", use_dummy_dataset=True)
+    
+    # Initialize RAG with GPT-Neo as the generator
+    rag_model = RagSequenceForGeneration.from_pretrained(
+        "facebook/rag-token-nq", retriever=retriever, generator=custom_generator
+    )
+    
+    return tokenizer, rag_model
 
-qa_pipeline = load_model()
+tokenizer, rag_model = load_gpt_neo_rag()
 
-def extract_text_from_pdf(pdf_file):
-    pdf_reader = PyPDF2.PdfReader(pdf_file)
+# Function to read resume PDF
+def read_pdf(file):
+    pdf_reader = PdfReader(file)
     text = ""
-    for page_num in range(len(pdf_reader.pages)):
-        page = pdf_reader.pages[page_num]
+    for page in pdf_reader.pages:
         text += page.extract_text()
     return text
 
-def generate_answer(question, context):
-    # Add context to the question
+# Function to generate a contextualized answer using RAG with GPT-Neo
+def generate_answer(question, resume_text, name="The candidate"):
+    """
+    Uses RAG with GPT-Neo to generate answers based on the resume.
+    """
+    # Add context instruction to guide the model
     context_instruction = (
-        f"Your professional experience is outlined in the following resume. "
+        f"You are {name}, and your professional experience is outlined in the following resume. "
         "Answer the question as if you are the candidate, providing details from the resume where relevant."
     )
-
+    
     # Combine the question with the context instruction
     full_question = f"{context_instruction} Question: {question}"
+    
+    # Tokenize the input
+    inputs = tokenizer(full_question, resume_text, return_tensors="pt", truncation=True, padding="longest")
+    
+    # Generate the response
+    outputs = rag_model.generate(**inputs)
+    
+    # Decode the generated response
+    answer = tokenizer.decode(outputs[0], skip_special_tokens=True)
+    
+    return answer
 
-    # Get the answer using DistilBERT pipeline
-    result = qa_pipeline(question=full_question, context=resume_text)
-
-# Streamlit app layout
-st.title("Resume-based Q&A Bot")
-
-st.write("Upload your resume as a PDF, and the bot will answer questions about your professional experiences.")
+# Streamlit app UI
+st.title("Resume-based Q&A Bot (RAG with GPT-Neo)")
 
-uploaded_file = st.file_uploader("Upload Resume (PDF)", type="pdf")
+st.write("Upload your resume and ask questions about your professional experience!")
 
-if uploaded_file:
-    st.write("Resume uploaded successfully!")
-    resume_text = extract_text_from_pdf(uploaded_file)
-    #st.text_area("Extracted Resume Text", value=resume_text, height=250)
+# File uploader for the resume
+uploaded_file = st.file_uploader("Upload your resume (PDF format)", type=["pdf"])
 
-    question = st.text_input("Ask a question about your professional experiences:")
+# If a file is uploaded, extract the text
+if uploaded_file is not None:
+    resume_text = read_pdf(uploaded_file)
+    st.write("Resume successfully uploaded!")
+    st.write("Extracted Resume Text:")
+    st.text(resume_text)  # Display the extracted resume text for reference
     
-    if st.button("Get Answer") and question:
-        with st.spinner("Generating answer..."):
-            answer = generate_answer(question, resume_text)
-            st.write("**Answer:**", answer)
-
+    # Text input for questions
+    question = st.text_input("Ask a question about the resume")
+    
+    # Name input for the person in the resume
+    candidate_name = st.text_input("Enter the candidate's name (optional)", "The candidate")
+    
+    # Generate and display the answer when the button is clicked
+    if st.button("Generate Answer"):
+        if question:
+            answer = generate_answer(question, resume_text, candidate_name)
+            st.write("Answer:")
+            st.write(answer)
+        else:
+            st.write("Please enter a question.")
+else:
+    st.write("Please upload a PDF resume to get started.")