Spaces:
Sleeping
Sleeping
File size: 3,687 Bytes
0c5ef1f 163b706 0c5ef1f 163b706 0c5ef1f 163b706 0c5ef1f 02a74ff |
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 |
import os
import torch
import gradio as gr
import pandas as pd
from sentence_transformers import SentenceTransformer, util
from transformers import GenerationConfig, AutoModelForCausalLM, AutoTokenizer
os.environ["GRADIO_ANALYTICS_ENABLED"] = "False"
class PersianRAG:
def __init__(self, knowledge, embedding_model='LABSE', llm_model="MehdiHosseiniMoghadam/AVA-Mistral-7B-V2",
device='cuda', retrieved_docs=3):
self.device = device
self.retrieved_docs = retrieved_docs
self.answer_df = (knowledge['Answer'])
self.embedder = SentenceTransformer(embedding_model)
self.question_embeddings = self.embedder.encode((knowledge['Question']), show_progress_bar=True,
convert_to_tensor=True)
self.model = AutoModelForCausalLM.from_pretrained(llm_model, torch_dtype=torch.float16, device_map="auto")
self.tokenizer = AutoTokenizer.from_pretrained(llm_model)
self.generation_config = GenerationConfig(
do_sample=True,
top_k=1,
temperature=0.99,
max_new_tokens=900,
pad_token_id=self.tokenizer.eos_token_id
)
def rag(self, query):
ans = {}
question_embedding = self.embedder.encode(query, convert_to_tensor=True)
hits = util.semantic_search(question_embedding, self.question_embeddings)
hits = hits[0]
for hit in hits[0:self.retrieved_docs]:
ans[hit['corpus_id']] = self.answer_df[hit['corpus_id']]
ans = pd.DataFrame(list(ans.items()), columns=['id', 'res'])
prompt = f'''
با توجه به شرایط زیر به این سوال پاسخ دهید:
{query},
متن نوشته:
{ans['res'][0]} - {ans['res'][1]} - {ans['res'][2]}
'''
prompt = f"### Human:{prompt}\n### Assistant:"
inputs = self.tokenizer(prompt, return_tensors="pt").to(self.device)
outputs = self.model.generate(**inputs, generation_config=self.generation_config)
return self.tokenizer.decode(outputs[0], skip_special_tokens=True)
# Function to load CSV and initialize PersianRAG
def init_rag(knowledge_file, embedding_model, llm_model, device, retrieved_docs):
knowledge = pd.read_csv(knowledge_file)
rag_system = PersianRAG(knowledge, embedding_model=embedding_model, llm_model=llm_model, device=device,
retrieved_docs=retrieved_docs)
return rag_system
# Function to handle querying
def query_rag(rag_system, query):
return rag_system.rag(query)
# Gradio interface to upload CSV and configure RAG system
def rag_interface(knowledge_file, query, embedding_model, llm_model, device, retrieved_docs):
rag_system = init_rag(knowledge_file, embedding_model, llm_model, device, retrieved_docs)
return query_rag(rag_system, query)
# Create Gradio interface
interface = gr.Interface(
fn=rag_interface,
inputs=[
gr.File(label="Upload Knowledge Base CSV"),
gr.Textbox(label="Enter your query"),
gr.Dropdown(choices=["LABSE", "paraphrase-multilingual-mpnet-base-v2"], value="LABSE", label="Embedding Model"),
gr.Textbox(value="MehdiHosseiniMoghadam/AVA-Mistral-7B-V2", label="LLM Model Name"),
gr.Dropdown(choices=["cuda", "cpu"], value="cuda", label="Device"),
gr.Slider(minimum=1, maximum=5, step=1, value=3, label="Number of Retrieved Documents")],
outputs="text",
title="Persian RAG System",
description="Upload a CSV file as the knowledge base, ask a question, and get an answer.")
# Launch the Gradio interface
interface.launch(debug=True)
|