|
import gradio as gr |
|
from sentence_transformers import SentenceTransformer |
|
from wikipediaapi import Wikipedia |
|
import textwrap |
|
import numpy as np |
|
from openai import OpenAI |
|
|
|
|
|
def process_query(wiki_page, model_name, embed_dim, query, api_key): |
|
model_mapping = { |
|
"Arabic-mpnet-base-all-nli-triplet": "Omartificial-Intelligence-Space/Arabic-mpnet-base-all-nli-triplet", |
|
"Arabic-all-nli-triplet-Matryoshka": "Omartificial-Intelligence-Space/Arabic-all-nli-triplet-Matryoshka", |
|
"Arabert-all-nli-triplet-Matryoshka": "Omartificial-Intelligence-Space/Arabert-all-nli-triplet-Matryoshka", |
|
"Arabic-labse-Matryoshka": "Omartificial-Intelligence-Space/Arabic-labse-Matryoshka", |
|
"Marbert-all-nli-triplet-Matryoshka": "Omartificial-Intelligence-Space/Marbert-all-nli-triplet-Matryoshka" |
|
} |
|
|
|
model_path = model_mapping[model_name] |
|
model = SentenceTransformer(model_path, trust_remote_code=True, truncate_dim=embed_dim) |
|
wiki = Wikipedia('RAGBot/0.0', 'ar') |
|
doc = wiki.page(wiki_page).text |
|
paragraphs = doc.split('\n\n') |
|
|
|
for i, p in enumerate(paragraphs): |
|
wrapped_text = textwrap.fill(p, width=100) |
|
|
|
docs_embed = model.encode(paragraphs, normalize_embeddings=True) |
|
query_embed = model.encode(query, normalize_embeddings=True) |
|
similarities = np.dot(docs_embed, query_embed.T) |
|
top_3_idx = np.argsort(similarities, axis=0)[-3:][::-1].tolist() |
|
most_similar_documents = [paragraphs[idx] for idx in top_3_idx] |
|
|
|
CONTEXT = "" |
|
for i, p in enumerate(most_similar_documents): |
|
wrapped_text = textwrap.fill(p, width=100) |
|
CONTEXT += wrapped_text + "\n\n" |
|
|
|
prompt = f""" |
|
use the following CONTEXT to answer the QUESTION at the end. |
|
If you don't know the answer, just say that you don't know, don't try to make up an answer. |
|
CONTEXT: {CONTEXT} |
|
QUESTION: {query} |
|
""" |
|
|
|
client = OpenAI(api_key=api_key) |
|
response = client.chat.completions.create( |
|
model="gpt-4o", |
|
messages=[ |
|
{"role": "user", "content": prompt}, |
|
] |
|
) |
|
|
|
return response.choices[0].message.content |
|
|
|
|
|
wiki_page_input = gr.Textbox(label="Wikipedia Page (in Arabic)") |
|
query_input = gr.Textbox(label="Query (in Arabic)") |
|
api_key_input = gr.Textbox(label="OpenAI API Key", type="password") |
|
|
|
model_choice = gr.Dropdown( |
|
choices=[ |
|
"Arabic-mpnet-base-all-nli-triplet", |
|
"Arabic-all-nli-triplet-Matryoshka", |
|
"Arabert-all-nli-triplet-Matryoshka", |
|
"Arabic-labse-Matryoshka", |
|
"Marbert-all-nli-triplet-Matryoshka" |
|
], |
|
label="Choose Embedding Model" |
|
) |
|
|
|
embed_dim_choice = gr.Dropdown( |
|
choices=[768, 512, 256, 128, 64], |
|
label="Embedding Dimension" |
|
) |
|
|
|
output_text = gr.Textbox(label="Output") |
|
|
|
gr.Interface( |
|
fn=process_query, |
|
inputs=[wiki_page_input, model_choice, embed_dim_choice, query_input, api_key_input], |
|
outputs=output_text, |
|
title="Arabic Wiki RAG", |
|
description="Choose a Wikipedia page, embedding model, and dimension to answer a query in Arabic." |
|
).launch() |
