Update app.py

app.py

@@ -1,93 +1,105 @@
 import gradio as gr
-from smolagents import CodeAgent, DuckDuckGoSearchTool
-from smolagents.models import TransformersModel
+from transformers import AutoTokenizer, AutoModelForCausalLM
+import torch
 
-class SimpleAIAgent:
+class FixedAIAgent:
     def __init__(self):
-        print("Initializing AI Agent...")
-        self.model = TransformersModel("microsoft/DialoGPT-small")
-
-        self.search_tool = DuckDuckGoSearchTool()
-
-        self.agent = CodeAgent(
-            tools=[self.search_tool],
-            model=self.model,
-            max_steps=4
-        )
-        
-        print("AI Agent ready!")
-
-    def chat(self, message, history):
-        """
-        Основная функция для общения с агентом
-        """
-        print(f"User asked: {message}")
-
-        prompt = f"""
-        The user asked: {message}
+        print("Loading DialoGPT-small...")
         
-        Please provide a helpful and accurate answer.
-        If you need current information, use the search tool to find it online.
-        Keep your response clear and conversational.
-        """
+        # Загружаем модель и токенизатор
+        self.tokenizer = AutoTokenizer.from_pretrained("microsoft/DialoGPT-small")
+        self.model = AutoModelForCausalLM.from_pretrained("microsoft/DialoGPT-small")
         
+        # Добавляем pad token если его нет
+        if self.tokenizer.pad_token is None:
+            self.tokenizer.pad_token = self.tokenizer.eos_token
+            
+        print("DialoGPT-small loaded successfully!")
+    
+    def chat(self, message, history):
         try:
-            response = self.agent.run(prompt)
- 
-            clean_response = self.clean_answer(response)
-            print(f"Agent replied: {clean_response[:100]}...")
+            # Форматируем историю для DialoGPT
+            input_text = self.format_conversation(message, history)
+            
+            # Токенизируем входной текст
+            inputs = self.tokenizer.encode(input_text + self.tokenizer.eos_token, return_tensors='pt')
+            
+            # Генерируем ответ
+            with torch.no_grad():
+                outputs = self.model.generate(
+                    inputs,
+                    max_length=1000,
+                    pad_token_id=self.tokenizer.eos_token_id,
+                    do_sample=True,
+                    temperature=0.7,
+                    top_k=50,
+                    top_p=0.95,
+                    repetition_penalty=1.2
+                )
             
-            return clean_response
+            # Декодируем ответ
+            response = self.tokenizer.decode(outputs[0], skip_special_tokens=True)
+            
+            # Извлекаем только новый ответ (убираем историю)
+            bot_response = self.extract_new_response(input_text, response)
+            
+            return bot_response
             
         except Exception as e:
-            error_msg = f"Sorry, I encountered an error: {str(e)}"
-            print(f"Error: {e}")
-            return error_msg
-
-    def clean_answer(self, answer):
-        """
-        Убираем техническую информацию из ответа агента
-        """
-        lines = answer.split('\n')
-        clean_lines = []
+            return f"Error: {str(e)}"
+    
+    def format_conversation(self, message, history):
+        """Форматирует историю чата для DialoGPT"""
+        # Начинаем с нового сообщения
+        conversation = f"User: {message}"
         
-        for line in lines:
-            lower_line = line.lower()
-            if any(word in lower_line for word in ['tool:', 'searching', 'step', 'using tool']):
-                continue
-                
-            if line.strip():
-                clean_lines.append(line)
-
-        result = '\n'.join(clean_lines).strip()
-
-        if len(result) > 1500:
-            result = result[:1497] + "..."
+        # Добавляем историю (последние 2-3 сообщения)
+        if history:
+            # Берем последние 2 обмена
+            recent_history = history[-2:] if len(history) > 2 else history
             
-        return result if result else "I couldn't find a good answer to that question."
-
-ai_agent = SimpleAIAgent()
+            for user_msg, bot_msg in recent_history:
+                conversation = f"User: {user_msg}\nBot: {bot_msg}\n" + conversation
+        
+        return conversation
+    
+    def extract_new_response(self, input_text, full_response):
+        """Извлекает только новый ответ из полного ответа модели"""
+        if input_text in full_response:
+            # Убираем входной текст чтобы оставить только новый ответ
+            new_response = full_response[len(input_text):].strip()
+            # Убираем возможные префиксы
+            if new_response.startswith("Bot:"):
+                new_response = new_response[4:].strip()
+            return new_response
+        else:
+            # Если не нашли входной текст, возвращаем как есть
+            return full_response
 
-with gr.Blocks(title="My AI Assistant") as chat_interface:
-    gr.Markdown("# My AI Assistant")
-    gr.Markdown("Ask me anything! I can search the internet for current information.")
+# Создаем агента
+agent = FixedAIAgent()
 
-    chatbot = gr.Chatbot(height=400)
+# Создаем интерфейс
+with gr.Blocks() as app:
+    gr.Markdown("# AI Chat Assistant")
+    gr.Markdown("Powered by DialoGPT-small")
+    
+    chatbot = gr.Chatbot(height=400, label="Chat History")
     msg = gr.Textbox(
-        label="Your question",
-        placeholder="Ask me anything...",
+        label="Your message", 
+        placeholder="Type your message here...",
         lines=2
     )
-    clear_btn = gr.Button("Clear Chat")
+    clear = gr.Button("Clear Chat")
     
     def respond(message, chat_history):
-        bot_response = ai_agent.chat(message, chat_history)
-        chat_history.append((message, bot_response))
+        bot_message = agent.chat(message, chat_history)
+        chat_history.append((message, bot_message))
         return "", chat_history
     
     msg.submit(respond, [msg, chatbot], [msg, chatbot])
-    clear_btn.click(lambda: None, None, chatbot, queue=False)
+    clear.click(lambda: None, None, chatbot)
 
 if __name__ == "__main__":
-    print("Starting AI Chat Assistant...")
-    chat_interface.launch(share=True)
+    print("Starting AI Chat...")
+    app.lalunch()