Upload App.py

App.py

@@ -0,0 +1,157 @@
+import gradio as gr
+import numpy as np
+from sklearn.metrics.pairwise import cosine_similarity
+# from sklearn.decomposition import PCA
+from langchain_community.llms import Ollama
+from langchain_chroma import Chroma
+import langchain
+from langchain_community.document_loaders import DirectoryLoader, TextLoader, PyPDFLoader
+
+from langchain_experimental.text_splitter import SemanticChunker
+from langchain.text_splitter import RecursiveCharacterTextSplitter
+from langchain_community.embeddings.ollama import OllamaEmbeddings
+
+from typing import List, Dict
+from langchain.docstore.document import Document
+
+from transformers import T5Tokenizer, T5ForConditionalGeneration
+
+tokenizer = T5Tokenizer.from_pretrained("Voicelab/vlt5-base-keywords")
+model = T5ForConditionalGeneration.from_pretrained("Voicelab/vlt5-base-keywords")
+
+vectorstore = Chroma(
+    # docs,
+    embedding_function=OllamaEmbeddings(model = "gemma:2b"),
+    persist_directory="chroma_db"
+)
+
+print(vectorstore.similarity_search_with_score("Course Leader"))
+
+llm = Ollama(
+    model="llama3.2:3b"
+)
+
+def retrieve_relevant_chunks(
+    vector_store: Chroma,
+    query: str,
+    n_docs: int = 2,
+    chunks_per_doc: int = 5
+) -> Dict[str, List[Document]]:
+
+    # Get more results initially to ensure we have enough unique documents
+    results = vector_store.similarity_search_with_score(
+        query,
+        k=50  # Fetch more to ensure we have enough unique documents
+    )
+    
+    # Group results by document ID
+    doc_chunks: Dict[str, List[tuple]] = {}
+    for doc, score in results:
+        doc_id = doc.metadata.get('source', '')  # or use appropriate metadata field
+        if doc_id:
+            if doc_id not in doc_chunks:
+                doc_chunks[doc_id] = []
+            doc_chunks[doc_id].append((doc, score))
+    
+    # Sort documents by their best matching chunk's score
+    sorted_docs = sorted(
+        doc_chunks.items(),
+        key=lambda x: min(chunk[1] for chunk in x[1])
+    )
+    
+    # Take only the top n_docs documents
+    top_docs = sorted_docs[:n_docs]
+    
+    # For each top document, get the best chunks_per_doc chunks
+    final_results: Dict[str, List[Document]] = {}
+    for doc_id, chunks in top_docs:
+        # Sort chunks by score (relevance)
+        sorted_chunks = sorted(chunks, key=lambda x: x[1])
+        # Take only the specified number of chunks and store just the Document objects
+        final_results[doc_id] = [chunk[0] for chunk in sorted_chunks[:chunks_per_doc]]
+    
+    return final_results
+
+def display_results(results: Dict[str, List[str]]) -> None:
+    """
+    Display the retrieved chunks in a formatted way.
+    
+    Args:
+        results: Dictionary mapping document IDs to lists of text chunks
+    """
+    prompt = " "
+    for doc_id, chunks in results.items():
+        # prompt += f"\nDocument ID: {doc_id}\n"
+        prompt += "-" * 50
+        for i, chunk in enumerate(chunks, 1):
+            # prompt += f"\nChunk {i}:"
+            prompt += str(chunk) + "\n"
+            # prompt += "-" * 30
+    return prompt
+
+def main(query):
+    
+    # Initialize your vector store (example)
+    # vector_store = Chroma(
+    #     persist_directory="path/to/your/vectorstore",
+    #     embedding_function=your_embedding_function
+    # )
+    
+    upd_query = "Keyword: " + query
+    input_ids = tokenizer.encode(upd_query, return_tensors="pt")
+    outputs = model.generate(input_ids)
+    output_sequence = tokenizer.decode(outputs[0], skip_special_tokens=True)
+    # print(output_sequence)
+    result_list = list(set(item.strip() for item in output_sequence.split(',')))
+    # print(result_list)
+    output_string = ", ".join(result_list)
+    print(output_string)
+
+    try:
+        results = retrieve_relevant_chunks(
+            vector_store=vectorstore,
+            query=output_string,
+            n_docs=2,
+            chunks_per_doc=5
+        )
+        
+        prompt = display_results(results)
+        
+    except Exception as e:
+        print(f"Error: {str(e)}")
+
+    formatted_prompt = f"""
+        You are an AI assistant. Your goal is to answer questions regarding student handbooks based on the following context provided. Make sure all the answers are within the given context:
+        {prompt}
+
+        Based on the above, answer the following question:
+        {query}
+        Give the answer in a clear and concise manner
+        """
+
+    response = llm.predict(formatted_prompt)
+
+    return response
+
+with gr.Blocks() as demo:
+    #gr.Image("../Documentation/Context Diagram.png", scale=2)
+    #gr(title="Your Interface Title")
+    gr.Markdown("""
+                <center> 
+                <span style='font-size: 50px; font-weight: Bold; font-family: "Graduate", serif'>                
+                IIT RAG Student Handbooks
+                </span>
+                </center>
+                """)
+    with gr.Group():
+        query = gr.Textbox(label="Question")
+        answer = gr.Textbox(label="Answer")
+
+    with gr.Row():
+        login_btn = gr.Button(value="Generate")
+    
+    login_btn.click(main, inputs=[query], outputs=answer)
+   
+# demo.launch(share = True, auth=authenticate)
+demo.launch(share = True)
+