a

.gitattributes

@@ -1,37 +1,37 @@
-*.7z filter=lfs diff=lfs merge=lfs -text
-*.arrow filter=lfs diff=lfs merge=lfs -text
-*.bin filter=lfs diff=lfs merge=lfs -text
-*.bz2 filter=lfs diff=lfs merge=lfs -text
-*.ckpt filter=lfs diff=lfs merge=lfs -text
-*.ftz filter=lfs diff=lfs merge=lfs -text
-*.gz filter=lfs diff=lfs merge=lfs -text
-*.h5 filter=lfs diff=lfs merge=lfs -text
-*.joblib filter=lfs diff=lfs merge=lfs -text
-*.lfs.* filter=lfs diff=lfs merge=lfs -text
-*.mlmodel filter=lfs diff=lfs merge=lfs -text
-*.model filter=lfs diff=lfs merge=lfs -text
-*.msgpack filter=lfs diff=lfs merge=lfs -text
-*.npy filter=lfs diff=lfs merge=lfs -text
-*.npz filter=lfs diff=lfs merge=lfs -text
-*.onnx filter=lfs diff=lfs merge=lfs -text
-*.ot filter=lfs diff=lfs merge=lfs -text
-*.parquet filter=lfs diff=lfs merge=lfs -text
-*.pb filter=lfs diff=lfs merge=lfs -text
-*.pickle filter=lfs diff=lfs merge=lfs -text
-*.pkl filter=lfs diff=lfs merge=lfs -text
-*.pt filter=lfs diff=lfs merge=lfs -text
-*.pth filter=lfs diff=lfs merge=lfs -text
-*.rar filter=lfs diff=lfs merge=lfs -text
-*.safetensors filter=lfs diff=lfs merge=lfs -text
-saved_model/**/* filter=lfs diff=lfs merge=lfs -text
-*.tar.* filter=lfs diff=lfs merge=lfs -text
-*.tar filter=lfs diff=lfs merge=lfs -text
-*.tflite filter=lfs diff=lfs merge=lfs -text
-*.tgz filter=lfs diff=lfs merge=lfs -text
-*.wasm filter=lfs diff=lfs merge=lfs -text
-*.xz 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
-db/default__vector_store.json filter=lfs diff=lfs merge=lfs -text
-db/docstore.json filter=lfs diff=lfs merge=lfs -text
+*.7z filter=lfs diff=lfs merge=lfs -text
+*.arrow filter=lfs diff=lfs merge=lfs -text
+*.bin filter=lfs diff=lfs merge=lfs -text
+*.bz2 filter=lfs diff=lfs merge=lfs -text
+*.ckpt filter=lfs diff=lfs merge=lfs -text
+*.ftz filter=lfs diff=lfs merge=lfs -text
+*.gz filter=lfs diff=lfs merge=lfs -text
+*.h5 filter=lfs diff=lfs merge=lfs -text
+*.joblib filter=lfs diff=lfs merge=lfs -text
+*.lfs.* filter=lfs diff=lfs merge=lfs -text
+*.mlmodel filter=lfs diff=lfs merge=lfs -text
+*.model filter=lfs diff=lfs merge=lfs -text
+*.msgpack filter=lfs diff=lfs merge=lfs -text
+*.npy filter=lfs diff=lfs merge=lfs -text
+*.npz filter=lfs diff=lfs merge=lfs -text
+*.onnx filter=lfs diff=lfs merge=lfs -text
+*.ot filter=lfs diff=lfs merge=lfs -text
+*.parquet filter=lfs diff=lfs merge=lfs -text
+*.pb filter=lfs diff=lfs merge=lfs -text
+*.pickle filter=lfs diff=lfs merge=lfs -text
+*.pkl filter=lfs diff=lfs merge=lfs -text
+*.pt filter=lfs diff=lfs merge=lfs -text
+*.pth filter=lfs diff=lfs merge=lfs -text
+*.rar filter=lfs diff=lfs merge=lfs -text
+*.safetensors filter=lfs diff=lfs merge=lfs -text
+saved_model/**/* filter=lfs diff=lfs merge=lfs -text
+*.tar.* filter=lfs diff=lfs merge=lfs -text
+*.tar filter=lfs diff=lfs merge=lfs -text
+*.tflite filter=lfs diff=lfs merge=lfs -text
+*.tgz filter=lfs diff=lfs merge=lfs -text
+*.wasm filter=lfs diff=lfs merge=lfs -text
+*.xz 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
+db/default__vector_store.json filter=lfs diff=lfs merge=lfs -text
+db/docstore.json filter=lfs diff=lfs merge=lfs -text

README.md

@@ -1,14 +1,14 @@
----
-title: Marco O1
-emoji: 💬
-colorFrom: yellow
-colorTo: purple
-sdk: gradio
-sdk_version: 5.7.1
-app_file: app.py
-pinned: false
-license: apache-2.0
-short_description: 'Marco-O1 A model similar to OpenAI o1 can reason  '
----
-
+---
+title: Marco O1
+emoji: 💬
+colorFrom: yellow
+colorTo: purple
+sdk: gradio
+sdk_version: 5.7.1
+app_file: app.py
+pinned: false
+license: apache-2.0
+short_description: 'Marco-O1 A model similar to OpenAI o1 can reason  '
+---
+
 An example chatbot using [Gradio](https://gradio.app), [`huggingface_hub`](https://huggingface.co/docs/huggingface_hub/v0.22.2/en/index), and the [Hugging Face Inference API](https://huggingface.co/docs/api-inference/index).

app.py

@@ -1,61 +1,87 @@
-import gradio as gr
-from transformers import pipeline
-
-# Загрузка модели Marco-o1 с квантизацией
-pipe = pipeline("text-generation", model="AIDC-AI/Marco-o1", device_map="auto", torch_dtype="auto", trust_remote_code=True)
-
-def respond(
-    message,
-    history: list[tuple[str, str]],
-    system_message,
-    max_tokens,
-    temperature,
-    top_p,
-):
-    messages = [system_message]
-    
-    for val in history:
-        if val[0]:
-            messages.append(val[0])
-        if val[1]:
-            messages.append(val[1])
-
-    messages.append(message)
-    
-    # Объединяем все сообщения в одну строку для передачи в модель
-    input_text = "\n".join(messages)
-    
-    response = pipe(
-        input_text,
-        max_length=max_tokens + len(input_text),
-        temperature=temperature,
-        top_p=top_p,
-        num_return_sequences=1
-    )[0]['generated_text']
-    
-    # Извлекаем новый ответ, исключая входные сообщения
-    new_response = response[len(input_text):].strip()
-    
-    yield new_response
-
-"""
-For information on how to customize the ChatInterface, peruse the gradio docs: https://www.gradio.app/docs/chatinterface
-"""
-demo = gr.ChatInterface(
-    respond,
-    additional_inputs=[
-        gr.Textbox(value="You are a friendly Chatbot.", label="System message"),
-        gr.Slider(minimum=1, maximum=2048, value=512, step=1, label="Max new tokens"),
-        gr.Slider(minimum=0.1, maximum=4.0, value=0.7, step=0.1, label="Temperature"),
-        gr.Slider(
-            minimum=0.1,
-            maximum=1.0,
-            value=0.95,
-            step=0.05,
-            label="Top-p (nucleus sampling)",
-        ),
-    ],
-)
-
-if __name__ == "__main__":
+import gradio as gr
+from transformers import pipeline
+from langchain.vectorstores import Chroma
+from langchain.embeddings import HuggingFaceEmbeddings
+import os
+
+# Load the embedding model
+embeddings = HuggingFaceEmbeddings(model_name="sentence-transformers/all-mpnet-base-v2")
+
+# Load the pre-existing vector database
+persist_directory = "db"
+vectordb = Chroma(persist_directory=persist_directory, embedding_function=embeddings)
+
+# Load the Marco-o1 model
+pipe = pipeline("text-generation", model="AIDC-AI/Marco-o1", device_map="auto", torch_dtype="auto", trust_remote_code=True)
+
+def get_relevant_context(query, k=3):
+    # Search the vector database for relevant documents
+    docs = vectordb.similarity_search(query, k=k)
+    # Combine the relevant documents into a single context string
+    context = "\n".join([doc.page_content for doc in docs])
+    return context
+
+def respond(
+    message,
+    history: list[tuple[str, str]],
+    system_message,
+    max_tokens,
+    temperature,
+    top_p,
+):
+    messages = [system_message]
+    
+    # Get relevant context from the vector database
+    context = get_relevant_context(message)
+    
+    # Add context to the system message
+    if context:
+        messages[0] = f"{system_message}\n\nRelevant context:\n{context}"
+    
+    for val in history:
+        if val[0]:
+            messages.append(val[0])
+        if val[1]:
+            messages.append(val[1])
+
+    messages.append(message)
+    
+    # Combine all messages into one string
+    input_text = "\n".join(messages)
+    
+    response = pipe(
+        input_text,
+        max_length=max_tokens + len(input_text),
+        temperature=temperature,
+        top_p=top_p,
+        num_return_sequences=1
+    )[0]['generated_text']
+    
+    # Extract new response
+    new_response = response[len(input_text):].strip()
+    
+    yield new_response
+
+demo = gr.ChatInterface(
+    respond,
+    additional_inputs=[
+        gr.Textbox(
+            value="You are a helpful AI assistant. Use the provided context to answer questions accurately.",
+            label="System message"
+        ),
+        gr.Slider(minimum=1, maximum=2048, value=512, step=1, label="Max new tokens"),
+        gr.Slider(minimum=0.1, maximum=4.0, value=0.7, step=0.1, label="Temperature"),
+        gr.Slider(
+            minimum=0.1,
+            maximum=1.0,
+            value=0.95,
+            step=0.05,
+            label="Top-p (nucleus sampling)",
+        ),
+    ],
+    title="Marco-O1 Assistant with Knowledge Base",
+    description="Ask questions about the documents in the knowledge base. The assistant will use the relevant context to provide accurate answers."
+)
+
+if __name__ == "__main__":
     demo.launch()

requirements.txt

@@ -1,4 +1,7 @@
-torch
-transformers
-bitsandbytes
-accelerate
+torch
+transformers
+bitsandbytes
+accelerate
+langchain
+chromadb
+sentence-transformers