Update app.py

app.py

@@ -28,9 +28,12 @@ def load_doc(list_file_path, chunk_size, chunk_overlap):
     text_splitter = RecursiveCharacterTextSplitter(chunk_size=chunk_size, chunk_overlap=chunk_overlap)
     return text_splitter.split_documents(pages)
 
-def create_db(splits, collection_name):
-    # Use the lightweight MiniLM model for embeddings
-    embedding = HuggingFaceEmbeddings(model_name="sentence-transformers/all-MiniLM-L6-v2")
+def create_db(splits, collection_name, db_type):
+    if db_type == 0:  # Multilingual MiniLM
+        embedding = HuggingFaceEmbeddings(model_name="sentence-transformers/paraphrase-multilingual-MiniLM-L12-v2")
+    else:  # Italian BERT
+        embedding = HuggingFaceEmbeddings(model_name="dbmdz/bert-base-italian-xxl-uncased")
+    
     new_client = chromadb.EphemeralClient()
     return Chroma.from_documents(documents=splits, embedding=embedding, client=new_client, collection_name=collection_name)
 
@@ -67,11 +70,11 @@ def create_collection_name(filepath):
         collection_name = collection_name[:-1] + 'Z'
     return collection_name
 
-def initialize_database(list_file_obj, chunk_size, chunk_overlap, progress=gr.Progress()):
+def initialize_database(list_file_obj, chunk_size, chunk_overlap, db_type, progress=gr.Progress()):
     list_file_path = [x.name for x in list_file_obj if x is not None]
     collection_name = create_collection_name(list_file_path[0])
     doc_splits = load_doc(list_file_path, chunk_size, chunk_overlap)
-    vector_db = create_db(doc_splits, collection_name)
+    vector_db = create_db(doc_splits, collection_name, db_type)
     return vector_db, collection_name, "Completed!"
 
 def initialize_LLM(llm_option, llm_temperature, max_tokens, top_k, vector_db, progress=gr.Progress()):
@@ -113,7 +116,7 @@ def demo():
             document = gr.Files(height=100, file_count="multiple", file_types=["pdf"], interactive=True, label="Carica i tuoi documenti PDF")
         
         with gr.Tab("Passo 2 - Elabora Documenti"):
-            db_btn = gr.Radio(["ChromaDB (MiniLM Embedding)"], label="Tipo di database vettoriale", value="ChromaDB (MiniLM Embedding)", type="index")
+            db_type = gr.Radio(["ChromaDB (Multilingual MiniLM Embedding)", "ChromaDB (Italian BERT Embedding)"], label="Tipo di database vettoriale", value="ChromaDB (Multilingual MiniLM Embedding)", type="index")
             with gr.Accordion("Opzioni Avanzate - Divisione del testo del documento", open=False):
                 slider_chunk_size = gr.Slider(100, 1000, 1000, step=20, label="Dimensione del chunk")
                 slider_chunk_overlap = gr.Slider(10, 200, 100, step=10, label="Sovrapposizione del chunk")
@@ -133,13 +136,13 @@ def demo():
         with gr.Tab("Passo 4 - Chatbot"):
             chatbot = gr.Chatbot(height=300)
             with gr.Accordion("Opzioni avanzate - Riferimenti ai documenti", open=False):
-                doc_sources = [gr.Textbox(label=f"Riferimento {i+1}", lines=2, container=True, scale=20) for i in range(5)]  # Modificato da range(3) a range(5)
-                source_pages = [gr.Number(label="Pagina", scale=1) for _ in range(5)]  # Modificato da range(3) a range(5)
+                doc_sources = [gr.Textbox(label=f"Riferimento {i+1}", lines=2, container=True, scale=20) for i in range(5)]
+                source_pages = [gr.Number(label="Pagina", scale=1) for _ in range(5)]
             msg = gr.Textbox(placeholder="Inserisci il messaggio (es. 'Di cosa tratta questo documento?')", container=True)
             submit_btn = gr.Button("Invia messaggio")
             clear_btn = gr.Button("Cancella conversazione")
            
-        db_btn.click(initialize_database, inputs=[document, slider_chunk_size, slider_chunk_overlap], outputs=[vector_db, collection_name, db_progress])
+        db_btn.click(initialize_database, inputs=[document, slider_chunk_size, slider_chunk_overlap, db_type], outputs=[vector_db, collection_name, db_progress])
         qachain_btn.click(initialize_LLM, inputs=[llm_btn, slider_temperature, slider_maxtokens, slider_topk, vector_db], outputs=[qa_chain, llm_progress])
 
         submit_btn.click(conversation, inputs=[qa_chain, msg, chatbot], outputs=[qa_chain, msg, chatbot] + doc_sources + source_pages)