Update agent.py

agent.py

@@ -9,7 +9,10 @@ from langgraph.graph import StateGraph, END
 from typing import Dict, Any
 from docx import Document
 from pptx import Presentation
-from langchain_ollama import ChatOllama
+#from langchain_ollama import ChatOllama
+#import ollama
+from transformers import AutoModelForCausalLM, AutoTokenizer, pipeline
+
 import logging
 import importlib.util
 import re
@@ -23,7 +26,7 @@ import torch
 from faster_whisper import WhisperModel
 from sentence_transformers import SentenceTransformer
 import faiss
-import ollama
+
 import asyncio
 #from shazamio import Shazam
 from langchain_community.document_loaders import WikipediaLoader, ArxivLoader
@@ -31,7 +34,7 @@ from bs4 import BeautifulSoup
 from typing import TypedDict, Optional
 from faiss import IndexFlatL2
 import pdfplumber
-from transformers import AutoModelForCausalLM, AutoTokenizer, pipeline
+
 from retrying import retry
 
 # Настройка путей для Hugging Face Spaces
@@ -122,11 +125,11 @@ def check_faiss():
         raise ImportError("faiss не установлена. Установите: pip install faiss-cpu")
     logger.info("faiss доступна.")
 
-def check_ollama():
-    if importlib.util.find_spec("ollama") is None:
-        logger.error("ollama не установлена. Установите: pip install ollama")
-        raise ImportError("ollama не установлена. Установите: pip install ollama")
-    logger.info("ollama доступна.")
+#def check_ollama():
+#    if importlib.util.find_spec("ollama") is None:
+#        logger.error("ollama не установлена. Установите: pip install ollama")
+#        raise ImportError("ollama не установлена. Установите: pip install ollama")
+#    logger.info("ollama доступна.")
 
 def check_shazamio():
     if importlib.util.find_spec("shazamio") is None:
@@ -143,18 +146,67 @@ def check_langchain_community():
 
 
 # Инициализация модели
+#try:
+#    llm = ChatOllama(base_url=OLLAMA_URL, model=MODEL_NAME, request_timeout=60)
+#    test_response = llm.invoke("Test")
+#    if test_response is None or not hasattr(test_response, 'content'):
+#        raise ValueError("Ollama модель недоступна или возвращает некорректный ответ")
+#    logger.info("Модель ChatOllama инициализирована.")
+#except Exception as e:
+#    logger.error(f"Ошибка инициализации модели: {e}")
+#    raise e
+
 try:
-    llm = ChatOllama(base_url=OLLAMA_URL, model=MODEL_NAME, request_timeout=60)
-    test_response = llm.invoke("Test")
-    if test_response is None or not hasattr(test_response, 'content'):
-        raise ValueError("Ollama модель недоступна или возвращает некорректный ответ")
-    logger.info("Модель ChatOllama инициализирована.")
+    device = "cuda" if torch.cuda.is_available() else "cpu"
+    logger.info(f"Используемое устройство: {device}")
+
+    # Инициализация Qwen2-7B
+    qwen_tokenizer = AutoTokenizer.from_pretrained("Qwen/Qwen2-7B-Instruct")
+    qwen_model = AutoModelForCausalLM.from_pretrained(
+        "Qwen/Qwen2-7B-Instruct",
+        device_map="auto",
+        load_in_4bit=True if device == "cuda" else False,  # Квантование для GPU
+        torch_dtype=torch.float16 if device == "cuda" else torch.float32
+    )
+    qwen_pipeline = pipeline(
+        "text-generation",
+        model=qwen_model,
+        tokenizer=qwen_tokenizer,
+        device_map="auto"
+    )
+    logger.info("Модель Qwen2-7B-Instruct инициализирована.")
+
+    # Инициализация Mixtral-8x7B
+    mixtral_tokenizer = AutoTokenizer.from_pretrained("mistralai/Mixtral-8x7B-Instruct-v0.1")
+    mixtral_model = AutoModelForCausalLM.from_pretrained(
+        "mistralai/Mixtral-8x7B-Instruct-v0.1",
+        device_map="auto",
+        load_in_4bit=True if device == "cuda" else False,
+        torch_dtype=torch.float16 if device == "cuda" else torch.float32
+    )
+    mixtral_pipeline = pipeline(
+        "text-generation",
+        model=mixtral_model,
+        tokenizer=mixtral_tokenizer,
+        device_map="auto"
+    )
+    logger.info("Модель Mixtral-8x7B-Instruct инициализирована.")
+
+    # Тестовый вызов для Qwen
+    test_input = qwen_tokenizer("Test", return_tensors="pt").to(device)
+    test_output = qwen_model.generate(**test_input, max_new_tokens=10)
+    test_response = qwen_tokenizer.decode(test_output[0], skip_special_tokens=True)
+    if not test_response:
+        raise ValueError("Qwen2-7B модель недоступна или возвращает пустой ответ")
+    logger.info(f"Тестовый ответ Qwen2-7B: {test_response}")
 except Exception as e:
-    logger.error(f"Ошибка инициализации модели: {e}")
+    logger.error(f"Ошибка инициализации моделей: {e}")
     raise e
 
 
 
+    
+
 
 # --- Состояние для LangGraph ---
 class AgentState(TypedDict):
@@ -421,7 +473,7 @@ async def process_file(file_path: str, question: str) -> str:
             check_faster_whisper()
             check_sentence_transformers()
             check_faiss()
-            check_ollama()
+            #check_ollama()
             transcribed_text = transcribe_audio(file_path)
             if transcribed_text.startswith("Error"):
                 logger.error(f"Ошибка транскрипции: {transcribed_text}")
@@ -844,7 +896,7 @@ def create_answer(state: AgentState) -> AgentState:
         file_content = state["file_content"]
         wiki_results = state["wiki_results"]
         arxiv_results = state["arxiv_results"]
-        web_results = state.get("web_results", None)  # Новое поле
+        web_results = state.get("web_results", None)
     except Exception as e:
         logger.error(f"Ошибка извлечения ключей: {str(e)}")
         return {"answer": f"Error extracting keys: {str(e)}", "raw_answer": f"Error extracting keys: {str(e)}"}
@@ -890,16 +942,15 @@ def create_answer(state: AgentState) -> AgentState:
     if "card game" in question_lower:
         logger.info("Обработка карточной игры...")
         cards = ["2 of clubs", "3 of hearts", "King of spades", "Queen of hearts", "Jack of clubs", "Ace of diamonds"]
-        # Шаги перестановок
-        cards = cards[3:] + cards[:3]  # 1. 3 карты сверху вниз
-        cards = [cards[1], cards[0]] + cards[2:]  # 2. Верхняя под вторую
-        cards = [cards[2]] + cards[:2] + cards[3:]  # 3. 2 карты сверху под третью
-        cards = [cards[-1]] + cards[:-1]  # 4. Нижняя наверх
-        cards = [cards[2]] + cards[:2] + cards[3:]  # 5. 2 карты сверху под третью
-        cards = cards[4:] + cards[:4]  # 6. 4 карты сверху вниз
-        cards = [cards[-1]] + cards[:-1]  # 7. Нижняя наверх
-        cards = cards[2:] + cards[:2]  # 8. 2 карты сверху вниз
-        cards = [cards[-1]] + cards[:-1]  # 9. Нижняя наверх
+        cards = cards[3:] + cards[:3]
+        cards = [cards[1], cards[0]] + cards[2:]
+        cards = [cards[2]] + cards[:2] + cards[3:]
+        cards = [cards[-1]] + cards[:-1]
+        cards = [cards[2]] + cards[:2] + cards[3:]
+        cards = cards[4:] + cards[:4]
+        cards = [cards[-1]] + cards[:-1]
+        cards = cards[2:] + cards[:2]
+        cards = [cards[-1]] + cards[:-1]
         state["answer"] = cards[0]
         state["raw_answer"] = cards[0]
         logger.info(f"Карточная игра обработана: {state['answer']}")
@@ -940,52 +991,20 @@ def create_answer(state: AgentState) -> AgentState:
             state["raw_answer"] = f"Error: {e}"
             return state
     
-
     # Обработка MP3-файлов
     file_path = state.get("file_path")
-
-
     if file_path and file_path.endswith(".mp3"):
         logger.info("Обработка MP3-файла")
-
-        if ext == ".mp3" and ("name of the song" in question.lower() or "what song" in question.lower()):
-            logger.warning("Распознавание песен больше не поддерживается: shazamio не установлена из-за конфликта с gradio. Но код работает в локальной версии без gradio")
-            return "Unknown"
+        if "name of the song" in question_lower or "what song" in question_lower():
+            logger.warning("Распознавание песен больше не поддерживается: shazamio не установлена")
+            state["answer"] = "Unknown"
+            state["raw_answer"] = "Song recognition not supported"
+            return state
         
-        #if "name of the song" in question_lower or "what song" in question_lower:
-        #    logger.info("Распознавание песни")
-        #    try:
-        #        # Поскольку file_content уже содержит результат process_file
-        #        if file_content and not file_content.startswith("Error"):
-        #            state["answer"] = file_content if file_content != "Not found" else "Unknown"
-        #            state["raw_answer"] = file_content
-        #            logger.info(f"Ответ для песни: {state['answer']}")
-        #        else:
-        #            state["answer"] = "Unknown"
-        #            state["raw_answer"] = "Error: No valid song recognition result"
-        #            logger.error("Ошибка: результат распознавания песни недоступен")
-        #        return state
-        #    except Exception as e:
-        #        logger.error(f"Ошибка распознавания песни: {str(e)}")
-        #        state["answer"] = "Unknown"
-        #        state["raw_answer"] = f"Error recognizing song: {str(e)}"
-        #        return state
-    
-
         if "how long" in question_lower and "minute" in question_lower:
             logger.info("Определение длительности аудио")
             try:
-                # audio_path = os.path.join(DATA_DIR, "test", file_path) if Path(
-                    # os.path.join(DATA_DIR, "test", file_path)).exists() else os.path.join(
-                    # DATA_DIR, "validation", file_path)
-                # if not Path(audio_path).exists():
-                    # logger.error(f"Аудиофайл не найден: {audio_path}")
-                    # state["answer"] = "Unknown"
-                    # state["raw_answer"] = "Error: Audio file not found"
-                    # return state
-                # audio = pydub.AudioSegment.from_file(audio_path)
                 audio = pydub.AudioSegment.from_file(file_path)
-                
                 duration_seconds = len(audio) / 1000
                 duration_minutes = round(duration_seconds / 60)
                 state["answer"] = str(duration_minutes)
@@ -997,6 +1016,7 @@ def create_answer(state: AgentState) -> AgentState:
                 state["answer"] = "Unknown"
                 state["raw_answer"] = f"Error: {e}"
                 return state
+        
         # RAG для MP3 (аудиокниги)
         logger.info("RAG-обработка для MP3 (аудиокниги)")
         try:
@@ -1009,7 +1029,6 @@ def create_answer(state: AgentState) -> AgentState:
             # Инициализация RAG
             check_sentence_transformers()
             check_faiss()
-            check_ollama()
             rag_model = SentenceTransformer("all-MiniLM-L6-v2")
             index, sentences, embeddings = create_rag_index(file_content, rag_model)
             question_embedding = rag_model.encode([question], convert_to_numpy=True)
@@ -1024,7 +1043,7 @@ def create_answer(state: AgentState) -> AgentState:
 
             # Промпт для MP3 с RAG
             prompt = (
-                "You are a highly precise assistant tasked with answering a question based solely on the provided context from an audiobook's transcribed text. "
+                "[INST] You are a highly precise assistant tasked with answering a question based solely on the provided context from an audiobook's transcribed text. "
                 "Do not use any external knowledge or assumptions beyond the context. "
                 "Extract the answer strictly from the context, ensuring it matches the question's requirements. "
                 "If the question asks for an address, return only the street number and name (e.g., '123 Main'), excluding city, state, or street types (e.g., Street, Boulevard). "
@@ -1034,31 +1053,26 @@ def create_answer(state: AgentState) -> AgentState:
                 "Provide only the final answer, without explanations or additional text.\n"
                 f"Question: {question}\n"
                 f"Context: {relevant_context}\n"
-                "Answer:"
+                "Answer: [/INST]"
             )
             logger.info(f"Промпт для RAG: {prompt[:200]}...")
 
-            # Вызов модели llama3:8b
-            response = ollama.generate(
-                model="llama3:8b",
-                prompt=prompt,
-                options={
-                    "num_predict": 100,
-                    "temperature": 0.0,
-                    "top_p": 0.9,
-                    "stop": ["\n"]
-                }
+            # Вызов Mixtral-8x7B
+            response = mixtral_pipeline(
+                prompt,
+                max_new_tokens=100,
+                temperature=0.0,
+                top_p=0.9,
+                do_sample=False,
+                return_full_text=False
             )
-            answer = response.get("response", "").strip() or "Not found"
-            logger.info(f"Ollama (llama3:8b) вернул ответ: {answer}")
+            answer = response[0]["generated_text"].strip() or "Not found"
+            logger.info(f"Mixtral-8x7B вернул ответ: {answer}")
 
             # Проверка адресов
             if "address" in question_lower:
-                # Удаляем типы улиц, город, штат
                 answer = re.sub(r'\b(St\.|Street|Blvd\.|Boulevard|Ave\.|Avenue|Rd\.|Road|Dr\.|Drive)\b', '', answer, flags=re.IGNORECASE)
-                # Удаляем город и штат (после запятых)
                 answer = re.sub(r',\s*[^,]+$', '', answer).strip()
-                # Убедимся, что остались только номер и имя улицы
                 match = re.match(r'^\d+\s+[A-Za-z\s]+$', answer)
                 if not match:
                     logger.warning(f"Некорректный формат адреса: {answer}")
@@ -1074,19 +1088,16 @@ def create_answer(state: AgentState) -> AgentState:
             state["raw_answer"] = f"Error RAG: {str(e)}"
             return state
 
-
-
-   
     # Обработка вопросов с изображениями и Википедией
     logger.info("Проверка вопросов с изображениями и Википедией")
     if file_path and file_path.endswith((".jpg", ".png")) and "wikipedia" in question_lower:
         logger.info("Обработка изображения с Википедией")
         if wiki_results and not wiki_results.startswith("Error"):
             prompt = (
-                f"Question: {question}\n"
+                f"[INST] Question: {question}\n"
                 f"Wikipedia Content: {wiki_results[:1000]}\n"
-                f"Instruction: Provide ONLY the final answer.\n"
-                "Answer:"
+                "Instruction: Provide ONLY the final answer.\n"
+                "Answer: [/INST]"
             )
             logger.info(f"Промпт для изображения с Википедией: {prompt[:200]}...")
         else:
@@ -1098,7 +1109,7 @@ def create_answer(state: AgentState) -> AgentState:
         # Общий случай
         logger.info("Обработка общего случая")
         prompt = (
-            f"Question: {question}\n"
+            f"[INST] Question: {question}\n"
             f"Instruction: Provide ONLY the final answer.\n"
             f"Examples:\n"
             f"- Number: '42'\n"
@@ -1114,7 +1125,7 @@ def create_answer(state: AgentState) -> AgentState:
             prompt += f"Wikipedia Results: {wiki_results[:1000]}\n"
             has_context = True
             logger.info(f"Добавлен wiki_results: {wiki_results[:50]}...")
-        if arxiv_results and not arxiv_results.startswith("Error"):
+        if arxiv_results and not wiki_results.startswith("Error"):
             prompt += f"Arxiv Results: {arxiv_results[:1000]}\n"
             has_context = True
             logger.info(f"Добавлен arxiv_results: {arxiv_results[:50]}...")
@@ -1128,20 +1139,21 @@ def create_answer(state: AgentState) -> AgentState:
             state["answer"] = "Unknown"
             state["raw_answer"] = "No context available"
             return state
-        prompt += "Answer:"
-        logger.info(f"Промпт для общего случая: {prompt[:200]}...")
+        prompt += "Answer: [/INST]"
+        logger.info(f"Промпт: {prompt[:200]}...")
     
-    # Вызов LLM (qwen2:7b для не-MP3 случаев)
-    logger.info("Вызов LLM")
+    # Вызов Qwen2-7B
+    logger.info("Вызов Qwen2-7B")
     try:
-        response = llm.invoke(prompt)
-        logger.info(f"Ответ от llm.invoke: {response}")
-        if response is None:
-            logger.error("llm.invoke вернул None")
-            state["answer"] = "Unknown"
-            state["raw_answer"] = "LLM response is None"
-            return state
-        raw_answer = getattr(response, 'content', str(response)).strip() or "Unknown"
+        response = qwen_pipeline(
+            prompt,
+            max_new_tokens=100,
+            temperature=0.0,
+            top_p=0.9,
+            do_sample=False,
+            return_full_text=False
+        )
+        raw_answer = response[0]["generated_text"].strip() or "Unknown"
         state["raw_answer"] = raw_answer
         logger.info(f"Raw answer: {raw_answer[:100]}...")
         
@@ -1151,59 +1163,24 @@ def create_answer(state: AgentState) -> AgentState:
         clean_answer = re.sub(r'[^\w\s.-]', '', clean_answer)
         logger.info(f"Clean answer: {clean_answer[:100]}...")
         
-        
-####################################################
-# Проверка на галлюцинации
-        # def is_valid_answer(question, answer, context):
-            # question_lower = question.lower()
-            # if "address" in question_lower:
-                # return bool(re.match(r'^\d+\s+[A-Za-z\s]+$', answer))
-            # if "how many" in question_lower or "number" in question_lower:
-                # return bool(re.match(r'^\d+(\.\d+)?$', answer))
-            # if "format" in question_lower and "A.B.C.D." in question:
-                # return bool(re.match(r'^[A-Z]\.[A-Z]\.[A-Z]\.[A-Z]\.', answer))
-            # if context and answer.lower() not in context.lower():
-                # return False
-            # return True
-
-        # if not is_valid_answer(question, clean_answer, file_content or wiki_results or web_results):
-            # logger.warning(f"Ответ не соответствует контексту: {clean_answer}")
-            # state["answer"] = "Unknown"
-            # state["raw_answer"] = "Invalid answer for context"
-            # return state
-
-        # # Энтропийная проверка (опционально)
-        # response = llm.invoke(prompt, return_logits=True)
-        # if response.logits:
-            # probs = np.exp(response.logits) / np.sum(np.exp(response.logits))
-            # entropy = -np.sum(probs * np.log(probs + 1e-10))
-            # if entropy > 2.0:
-                # logger.warning(f"Высокая энтропия ответа: {entropy}")
-                # state["answer"] = "Unknown"
-                # state["raw_answer"] = "High uncertainty in response"
-                # return state 
-####################################################     
-        
-       
-        
         if any(keyword in question_lower for keyword in ["how many", "number", "score", "difference", "citations"]):
             match = re.search(r"\d+(\.\d+)?", clean_answer)
-            state["answer"] = match.group(0) if match else "Unknown"
+            state["answer"] = match.group(0) if match else "Not found"
         elif "stock price" in question_lower:
             match = re.search(r"\d+\.\d+", clean_answer)
-            state["answer"] = match.group(0) if match else "Unknown"
+            state["answer"] = match.group(0) if match else "Not found"
         elif any(keyword in question_lower for keyword in ["name", "what is", "restaurant", "city", "replica", "line", "song"]):
-            state["answer"] = clean_answer.split("\n")[0].strip() or "Unknown"
+            state["answer"] = clean_answer.split("\n")[0].strip() or "Not found"
         elif "address" in question_lower:
             match = re.search(r"\d+\s+[A-Za-z\s]+", clean_answer)
-            state["answer"] = match.group(0) if match else "Unknown"
+            state["answer"] = match.group(0) if match else "Not found"
         elif "The adventurer died" in clean_answer:
             state["answer"] = "The adventurer died."
         elif any(keyword in question_lower for keyword in ["code", "identifier", "issn"]):
             match = re.search(r"[\w-]+", clean_answer)
-            state["answer"] = match.group(0) if match else "Unknown"
+            state["answer"] = match.group(0) if match else "Not found"
         else:
-            state["answer"] = clean_answer.split("\n")[0].strip() or "Unknown"
+            state["answer"] = clean_answer.split("\n")[0].strip() or "Not found"
         
         logger.info(f"Final answer: {state['answer'][:50]}...")
         logger.info(f"Сгенерирован ответ: {state['answer'][:50]}...")
@@ -1214,7 +1191,6 @@ def create_answer(state: AgentState) -> AgentState:
     
     return state
 
-
  
    
 # --- Создание графа ---