Upload 15 files

.gitattributes

@@ -33,3 +33,5 @@ saved_model/**/* filter=lfs diff=lfs merge=lfs -text
 *.zip filter=lfs diff=lfs merge=lfs -text
 *.zst filter=lfs diff=lfs merge=lfs -text
 *tfevents* filter=lfs diff=lfs merge=lfs -text
+data/chunks.json filter=lfs diff=lfs merge=lfs -text
+data/embeddings.index filter=lfs diff=lfs merge=lfs -text

app.py

@@ -0,0 +1,27 @@
+import gradio as gr
+from src.retriever import CustomRetriever
+from src.generator import CustomGenerator
+from src.model import MyRagModel
+
+retriever = CustomRetriever(r'data/chunks.json', r'data/embeddings.index', r'data/metadata.json')
+generator = CustomGenerator()
+rag = MyRagModel(retriever, generator)
+
+
+def generate_response(query):
+    return rag.get_answer(query)
+
+
+interface = gr.Interface(
+    fn=generate_response,
+    inputs=gr.Textbox(label='Введите ваш вопрос: '),
+    outputs=gr.Textbox(label='Ответ модели: '),
+    title='Rag Model Demo',
+    description='QA ассистент по вопросам законов Кыргызской Республики. Может отвечать на вопросы по типу: '
+                '"Сколько лет дают за убийство человека?"',
+    allow_flagging='never',
+)
+
+
+if __name__ == '__main__':
+    interface.launch(share=True)

data/chunks.json

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:5490a3392697cd1e719445197aa3ebf4bc1ee6e3b954f231c1cc1ba97e740b61
+size 19871573

data/embeddings.index

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:feb52c93f3dbfe55c29a6fd605898862ea2fe483687451f0706c07fa06984118
+size 46215213

src/generator.py

@@ -0,0 +1,35 @@
+from groq import Client
+
+
+class CustomGenerator:
+    def __init__(self, model_name='llama-3.3-70b-versatile',
+                 api_key='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'):
+        self.model = model_name
+        self.client = Client(api_key=api_key)
+        self.system_prompt = """
+        Ты Q&A ассистент по вопросам законов Кыргызской Республики (Кыргызстана).
+        Твоя задача - ответить на заданный вопрос на основе переданной тебе информации.
+        Если вопрос касается законов Кыргызской республики, но в контексте нет ответа,
+        то ответь, что ты не знаешь ответа на этот вопрос. А если же вопрос не касается
+        законов Кыргызской Республики, то ответь только то, что ты не знаешь, потому что
+        ты помощник по вопросам законов Кыргызской Республики!
+        Укажи исходя из какой книги ответ, и сошлись на номер статьи, в конце ответа приведи добавь ссылку на книгу.
+        Ответ обязательно должен быть на русском, если ты придумал его на другом языке, то потом обязательно переведи на русский.
+        """
+
+    def generate(self, query, context):
+        query_wrapper = f"""
+        На основе вот этой информации {context},
+        ответь на вопрос клиента {query},
+        не забудь добавить источник ответа и
+        не забудь, что ответ должен быть на русском.
+        """
+        response = self.client.chat.completions.create(
+            model=self.model,
+            messages=[
+                {"role": "system", "content": self.system_prompt},
+                {"role": "user", "content": query_wrapper},
+            ]
+        )
+
+        return response.choices[0].message.content

src/model.py

@@ -0,0 +1,13 @@
+from .retriever import CustomRetriever
+from .generator import CustomGenerator
+
+
+class MyRagModel:
+    def __init__(self, retriever: CustomRetriever, generator: CustomGenerator):
+        self.retriever = retriever
+        self.generator = generator
+
+    def get_answer(self, query):
+        context = self.retriever.retrieve(query)
+        response = self.generator.generate(query, context)
+        return response

src/retriever.py

@@ -0,0 +1,30 @@
+from sentence_transformers import CrossEncoder
+from sentence_transformers import SentenceTransformer
+import faiss
+import json
+
+
+class CustomRetriever:
+    def __init__(self, chunks_path, embeddings_path, metadata_path, top_k=50):
+        self.model_bi = SentenceTransformer("deepvk/USER-bge-m3")
+        self.model_cross = CrossEncoder("DiTy/cross-encoder-russian-msmarco")
+        with open(chunks_path, "r") as f:
+            self.chunks = json.load(f)
+        self.index = faiss.read_index(embeddings_path)
+        self.top_k = top_k
+        with open(metadata_path, "r") as f:
+            self.metadata = json.load(f)
+
+    def retrieve(self, query):
+        query_vector = self.model_bi.encode([query])
+        faiss.normalize_L2(query_vector)
+        distances, indices = self.index.search(query_vector, self.top_k)
+        possible_answers = list()
+        for i in range(len(indices[0])):
+            possible_answers.append(self.chunks[indices[0][i]])
+        s = self.model_cross.rank(query, possible_answers)
+        context = ''
+        for i in range(5):
+            meta = self.metadata[str(indices[0][s[i]["corpus_id"]])]
+            context += f"Факт {str(i + 1)}: {possible_answers[s[i]['corpus_id']]}. Источник:\nкнига - {meta['book']}\nномер статьи - {meta['article_num']}\nссылка на книгу - {meta['link']}\n"
+        return context