app.py

# -*- coding: utf-8 -*-
"""Копія записника "Копія записника "Untitled8.ipynb""

Automatically generated by Colab.

Original file is located at
    https://colab.research.google.com/drive/1MdqEwbEry_Z-wnu7pkXI_bvJD0Qz26VU
"""


from datasets import load_dataset
from sentence_transformers import SentenceTransformer, util
import torch
from llama_index.core.node_parser import SentenceSplitter
from llama_index.core import Document
from rank_bm25 import BM25Okapi
from nltk.tokenize import word_tokenize
import gradio as gr
from llama_index.core.schema import NodeRelationship
from groq import Groq

import nltk

ds = load_dataset("lucadiliello/english_wikipedia", split='train[:5000]')

dataset=ds




arr=[Document(id_=i["url"], text=i["maintext"]) for i in dataset]
#documents = Document(id_=dataset["url"], text=dataset["maintext"])
splitter = SentenceSplitter(
    chunk_size=524,
    chunk_overlap=20,
)
nodes = splitter.get_nodes_from_documents(arr)

"""### BM25 & Minilm"""

nltk.download('punkt_tab')



# Prepare corpus for BM25
corpus = [node.text for node in nodes]
tokenized_corpus = [word_tokenize(doc.lower()) for doc in corpus]

# Initialize BM25
bm25 = BM25Okapi(tokenized_corpus)

# Define BM25 search function
def search_bm25(query, top_k=5):
    tokenized_query = word_tokenize(query.lower())
    scores = bm25.get_scores(tokenized_query)
    ranked_indices = sorted(range(len(scores)), key=lambda i: scores[i], reverse=True)
    return [(nodes[i], scores[i]) for i in ranked_indices[:top_k]]




corpus_embeddings=torch.load("./my_model (1)")
model = SentenceTransformer('all-MiniLM-L6-v2')

# Define semantic search function
def search_semantic(query, top_k=None):
    """
    Perform semantic search to find relevant documents.

    Args:
        query (str): The search query.
        top_k (int or None): Number of top results to return. If None, return scores for all documents.

    Returns:
        list: A list of tuples (Document, score) sorted by relevance.
    """
    # Encode the query
    query_embedding = model.encode(query, convert_to_tensor=True)

    # Compute similarity scores
    scores = util.cos_sim(query_embedding, corpus_embeddings)[0]

    # Handle top_k=None
    if top_k is None or top_k > len(scores):
        top_k = len(scores)

    # Get the top_k results
    ranked_indices = torch.topk(scores, k=top_k).indices
    return [(nodes[i], scores[i].item()) for i in ranked_indices]

def combined_search(query, bm25_weight=0.5, semantic_weight=0.5, top_k=5):
    # BM25 results
    bm25_results = search_bm25(query, top_k=None)  # Get scores for all documents

    # Semantic results
    semantic_results = search_semantic(query, top_k=None)  # Get scores for all documents

    # Normalize scores
    bm25_scores = {doc.doc_id: score for doc, score in bm25_results}
    semantic_scores = {doc.doc_id: score for doc, score in semantic_results}

    # Combine scores
    combined_scores = {}
    for doc in nodes:
        bm25_score = bm25_scores.get(doc.doc_id, 0)
        semantic_score = semantic_scores.get(doc.doc_id, 0)
        combined_score = bm25_weight * bm25_score + semantic_weight * semantic_score
        combined_scores[doc.doc_id] = combined_score

    # Sort by combined score
    sorted_docs = sorted(combined_scores.items(), key=lambda x: x[1], reverse=True)

    # Retrieve top_k results
    return [(next(node for node in nodes if node.doc_id == doc_id), score) for doc_id, score in sorted_docs[:top_k]]



def search(query, method="bm25", top_k=5, bm25_weight=0.5, semantic_weight=0.5):
    """
    Search documents using BM25, semantic search, or a combined method.

    Args:
        query (str): The search query.
        method (str): The retrieval method: "bm25", "semantic", or "combined".
        top_k (int): Number of top results to return.
        bm25_weight (float): Weight for BM25 in combined search (default: 0.5).
        semantic_weight (float): Weight for semantic search in combined search (default: 0.5).

    Returns:
        list: A list of tuples containing (Document, score).
    """
    if method == "bm25":
        return search_bm25(query, top_k=top_k)
    elif method == "semantic":
        return search_semantic(query, top_k=top_k)
    elif method == "combined":
        # Get scores from both BM25 and semantic search
        bm25_results = search_bm25(query, top_k=None)  # Retrieve scores for all documents
        semantic_results = search_semantic(query, top_k=None)  # Retrieve scores for all documents

        # Normalize and combine scores
        bm25_scores = {doc.id_: score for doc, score in bm25_results}
        semantic_scores = {doc.id_: score for doc, score in semantic_results}

        combined_scores = {}
        for doc in nodes:
            bm25_score = bm25_scores.get(doc.id_, 0)
            semantic_score = semantic_scores.get(doc.id_, 0)
            combined_score = bm25_weight * bm25_score + semantic_weight * semantic_score
            combined_scores[doc.id_] = combined_score

        # Sort by combined score
        sorted_docs = sorted(combined_scores.items(), key=lambda x: x[1], reverse=True)

        # Retrieve top_k results
        return [(next(node for node in nodes if node.id_ == id_), score) for id_, score in sorted_docs[:top_k]]
    else:
        raise ValueError("Invalid search method. Choose 'bm25', 'semantic', or 'combined'.")


def call_llm(documents, question, api_key):
  # Питання, яке має бути вставлене в промпт question = "Where is George Town airport"

  # Формуємо промпт з автоматично вставленими документами та питанням
  prompt = f"""
  Give me a concise answer to the following question: "{question}".
  Use only the following documents:
  """
  retrieved_docs = ""
  # Додаємо кожен документ з посиланням
  for iter,doc in enumerate(documents):
    link=None

    link = doc.relationships[NodeRelationship.SOURCE].node_id



    retrieved_docs +=str(iter+1)+"."+doc.text +"\n"*3+"_"*100+"\n"*3
    prompt += f'```{doc.text}``` Link:```{link}```\n'

  # Завершуємо промпт з інструкцією для відповіді
  prompt += """
  In your response, specify which part of the answer was obtained from which document, marking them as [1], [2], etc. After your answer, provide a list of the links with captions, but only include the links that were actually used in your answer.

  Use this for separating Links and Text:`--- Links Below ---`
  In link section, write links like: [number of link] link. Do not write anything else in this section.
  Explanation for LLM (Language Learning Model):
  - The text in quotes, for example "```doc```" and "Link:```{link}```", is used to represent the names of documents and their corresponding URLs or links.
  - In this case, "Link:```{link}```" refers to the title or name of a document, and "Link:```{link}```" refers to the URL or web address where the document can be found.
  - Your task is to refer to the specific document and the corresponding link when generating the answer. For instance, if part of the answer is from "Document 1", you should indicate that in your response as [1].
  - After providing the answer, include only the links that you actually used in your response, and list them with captions. If you did not use a particular document, do not include its link.
  """



  print(prompt)
  client = Groq(
      api_key=api_key,
  )

  chat_completion = client.chat.completions.create(
      messages=[
          {
              "role": "user",
              "content": prompt[:],
          }
      ],
      model="llama3-70b-8192",
  )
  response=chat_completion.choices[0].message.content

  response_splitted = response.split("--- Links Below ---")
  if len(response_splitted) > 1:
    main_answer = response_splitted[0]
    links_section = response_splitted[1]
  else:
    main_answer = response_splitted[0]
    links_section = ""
  print(links_section.strip())





  return main_answer.strip(), retrieved_docs, links_section.strip()
# Тепер у вас є main_answer для основної частини та links_section для лінків
#print("Main Answer:", main_answer.strip())
#print("Links Section:", links_section.strip())



# Function to process user input
def process_input(api_key, query, search_type, keyword_percentage, num_docs):
  try:
    if search_type == "Full Search" and (keyword_percentage < 0 or keyword_percentage > 100):
        return "Invalid percentage. Please enter a value between 0 and 100.", [], [], ""
    if num_docs <= 0:
        return "Number of documents must be greater than 0.", [], [], ""

    # Simulating search methods
    retrieved_docs = []
    sources = []

    if search_type == "Keyword Search":
        # BM25
        results = search(query, method="bm25", top_k=num_docs)

    elif search_type == "Semantic Search":
        # Semantic
        results = search(query, method="semantic",top_k=num_docs)
    elif search_type == "Full Search":
        # Split documents based on percentage
        keyword_count = int(keyword_percentage / 100 * num_docs)
        semantic_count = num_docs - keyword_count

        results = search(query, method="combined", bm25_weight=keyword_percentage/100, semantic_weight=(100-keyword_percentage)/100, top_k=num_docs)

    docs=([i[0] for i in results])
    print(docs)

    response, retrieved_docs, sources_text = call_llm(docs, query, api_key)
  except Exception as e:
    sources_text=e
    print(e)
    retrieved_docs=""
    response=""

  return response, retrieved_docs, sources_text

# Service description
description = """
## Retrieval-augmented generation by **Vitalii Pikhotskii** 😎 and **Sviatoslav Shainoha** 😉

Enter your LLM access key and a query in the text fields below. Use the checkboxes to enable/disable different search methods:
1. **Keyword Search:** Searches by keywords in documents (BM25).
2. **Semantic Search:** Uses semantic similarity for document retrieval.
3. **Full Search:** Combines both methods. Specify the percentage of keyword and semantic retrieval.
"""

with gr.Blocks() as demo:
    gr.Markdown(description)

    with gr.Row():
        api_key_input = gr.Textbox(label="LLM API Key", placeholder="Enter your key")
        query_input = gr.Textbox(label="Query", placeholder="For example: Who is Marry Cassatt?")

    with gr.Row():
        search_type_selector = gr.Radio(
            choices=["Keyword Search", "Semantic Search", "Full Search"],
            label="Search Type",
            value="Keyword Search"
        )

    with gr.Row(visible=False) as full_search_options:
        keyword_percentage_slider = gr.Slider(
            label="Keyword Search Percentage",
            minimum=0,
            maximum=100,
            step=1,
            value=50
        )

    search_type_selector.change(
        lambda choice: gr.update(visible=choice == "Full Search"),
        inputs=search_type_selector,
        outputs=full_search_options
    )

    num_docs_input = gr.Number(
        label="Number of Documents",
        value=5,
        precision=0,
        minimum=1
    )

    with gr.Row():
        submit_btn = gr.Button("Submit")

    with gr.Row():
        response_output = gr.Textbox(label="System Response")


    with gr.Row():
        sources_output = gr.Textbox(label="Sources", lines=5)

    with gr.Row():
        retrieved_docs_output = gr.Textbox(label="Documents (Retriever)", lines=5)
        #reranked_docs_output = gr.Textbox(label="Documents (Reranker)", lines=5)

    submit_btn.click(
        process_input,
        inputs=[
            api_key_input,
            query_input,
            search_type_selector,
            keyword_percentage_slider,
            num_docs_input
        ],
        outputs=[response_output, retrieved_docs_output, sources_output]
    )

# Launching the app
demo.launch()