Update app.py

app.py

@@ -60,24 +60,29 @@ def initialize_database(document, chunk_size, chunk_overlap, progress=gr.Progres
     logger.info("Database initialized successfully.")
     return vectorstore, None, "Initialized"  # Aggiunto None come secondo output
 
-def initialize_LLM(llm_option, llm_temperature, max_tokens, top_k, vector_db, progress=gr.Progress(), language="italian"):
+def initialize_LLM(llm_option, llm_temperature, max_tokens, top_k, vector_db, progress=gr.Progress(), language="italiano"):
     logger.info("Initializing LLM chain...")
-    
+
     # Define the default LLMS based on the language
     if language == "italiano":
         default_llm = "google/gemma-7b-it"
     else:
         default_llm = "mistralai/Mistral-7B-Instruct-v0.2"
 
-    # Create an instance of the LLM
+    # Try to load the tokenizer and model
     try:
-        llm = LLM.from_pretrained(default_llm)
+        tokenizer = AutoTokenizer.from_pretrained(default_llm)
+        model = AutoModelForCausalLM.from_pretrained(default_llm)
     except Exception as e:
         logger.error(f"Error initializing LLM: {e}")
         return None, "Failed to initialize LLM"
 
+    # Resize token embeddings if needed
+    if len(tokenizer) > model.config.max_position_embeddings:
+        model.resize_token_embeddings(len(tokenizer))
+
     qa_chain = ConversationalRetrievalChain.from_llm(
-        llm=llm,
+        llm=model,
         retriever=vector_db.as_retriever(),
         chain_type="stuff",
         temperature=llm_temperature,