File size: 5,804 Bytes
2ee953c
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
import streamlit as st
from langchain_pinecone import PineconeVectorStore
from langchain_openai import OpenAI, ChatOpenAI
from langchain_google_genai import ChatGoogleGenerativeAI
from langchain.prompts import PromptTemplate
from langchain_huggingface import HuggingFaceEmbeddings
from langchain.output_parsers import PydanticOutputParser
from dotenv import load_dotenv
import os
from pydantic import BaseModel, Field
from typing import List, Union
from typing_extensions import Literal


# Pydantic Schema
class Question(BaseModel):
    question: str = Field(..., description="The question prompt that the user needs to answer.")
    type: Literal['fill_missing', 'MCQ', 'short_answer'] = Field(..., description="The type of question: fill_missing, MCQ, or short_answer.")
    options: Union[List[str], None] = Field(None, description="The options for the question, used only for MCQ type.")

class Questions(BaseModel):
    no_of_questions: int = Field(..., description="The total number of questions generated.")
    questions: List[Question] = Field(..., description="A list of Question objects.")


# Function to download Hugging Face embeddings
def download_hugging_face_embeddings():
    embeddings = HuggingFaceEmbeddings(model_name='sentence-transformers/all-MiniLM-L6-v2')
    return embeddings


# Load environment variables
def load_env_variables():
    load_dotenv()
    PINECONE_API_KEY = os.environ.get('PINECONE_API_KEY')
    OPENAI_API_KEY = os.environ.get('OPENAI_API_KEY')
    GOOGLE_API_KEY = os.environ.get("GOOGLE_API_KEY")

    os.environ["PINECONE_API_KEY"] = PINECONE_API_KEY
    os.environ["OPENAI_API_KEY"] = OPENAI_API_KEY
    os.environ["GOOGLE_API_KEY"] = GOOGLE_API_KEY


# Function to initialize Pinecone vector store and retriever
def initialize_vector_store(embeddings):
    index_name = "yolotest"
    vector_store = PineconeVectorStore.from_existing_index(index_name=index_name, embedding=embeddings)
    retriever = vector_store.as_retriever(search_type="similarity", search_kwargs={"k": 5})
    return retriever


# Function to initialize LLM
def initialize_llm():
    llm = ChatOpenAI(api_key=os.environ.get("OPENAI_API_KEY"), temperature=0, model='gpt-3.5-turbo-0125')
    llm_structured = llm.with_structured_output(Questions)
    return llm_structured


# Function to initialize prompt template
def initialize_prompt_template(parser):
    prompt_template = """

    You are given some context below. Based on the context, generate questions with one of the following types: ['fill_missing', 'MCQ', 'short_answer'].



    Context: {context}



    Make sure to generate a variety of questions. The types should be distributed among 'fill_missing', 'MCQ', and 'short_answer'. 



    For MCQs, provide at least 3 options. For 'fill_missing', make sure to leave a gap that can be filled. For 'short_answer', make sure the answer is clear from the context.



    Return a list of questions, with each question having the following structure:

    - 'question': The question prompt.

    - 'type': The type of question: 'fill_missing', 'MCQ', or 'short_answer'.

    - 'options': For 'MCQ', a list of options. For other types, this field should be omitted.



    The response should be with a format that matches the following structure:

    {format_instructions}

    """
    return prompt_template


# Initialize components before user query
embeddings = download_hugging_face_embeddings()
load_env_variables()
retriever = initialize_vector_store(embeddings)
llm_structured = initialize_llm()
parser = PydanticOutputParser(pydantic_object=Questions)
prompt_template = initialize_prompt_template(parser)


# Function to generate questions from the context
def generate_questions_from_context(query: str, retriever, llm_structured, prompt_template, parser):
    # Retrieve relevant documents from the vector store using the retriever
    retrieved_docs = retriever.invoke(query)
    retrieved_docs = [rec.page_content for rec in retrieved_docs]

    # Combine the retrieved documents into a single context string
    context = " ".join([doc for doc in retrieved_docs])

    # Initialize prompt template
    prompt = PromptTemplate(template=prompt_template, input_variables=["context"],
                            partial_variables={"format_instructions": parser.get_format_instructions()})

    # Create chain and generate response
    chain = prompt | llm_structured
    response = chain.invoke({
        "context": context, "format_instructions": parser.get_format_instructions()
    })
    return response


# Streamlit Interface
def main():
    st.title("A Simple RAG App to Generate Questions in Specific Formats")

    # Display user query input
    query = st.chat_input("Say something: ")
    if query:
        with st.spinner('Generating questions...'):
            st.write(f"Your query: {query}")  # Display user query
            
            # Generate questions
            result = generate_questions_from_context(query, retriever, llm_structured, prompt_template, parser)
            st.write(result)
            # # Display the result in a more readable format
            # st.subheader("Generated Questions")
            # for i, question in enumerate(result.questions):
            #     st.markdown(f"### Question {i + 1}: {question.question}")
            #     st.markdown(f"**Type**: {question.type}")
            #     if question.type == "MCQ" and question.options:
            #         st.markdown("**Options**:")
            #         for option in question.options:
            #             st.markdown(f"- {option}")
            #     st.markdown("---")


# Run the app
if __name__ == "__main__":
    main()