Update app.py

app.py

@@ -8,9 +8,7 @@ from langdetect import detect
 # Проверяем наличие текстовых файлов и читаем их
 def load_text_files():
     files = {
-        "vampires": "vampires.txt",
-        "werewolves": "werewolves.txt",
-        "humans": "humans.txt"
+        "vampires": "vampires.txt"
     }
     
     loaded_data = {}
@@ -66,22 +64,16 @@ def create_knowledge_base(text_data, embed_fn):
     
     return collection
 
-# Инициализация модели для ответов
+# Инициализация модели для ответов (упрощенная версия)
 def initialize_llm_model():
-    from transformers import AutoModelForCausalLM, AutoTokenizer, pipeline
+    from transformers import pipeline
     
-    model_name = "IlyaGusev/saiga_mistral_7b"
-    tokenizer = AutoTokenizer.from_pretrained(model_name, use_fast=False)
-    model = AutoModelForCausalLM.from_pretrained(model_name)
-    
-    pipe = pipeline(
+    # Используем меньшую модель для Hugging Face Spaces
+    return pipeline(
         "text-generation",
-        model=model,
-        tokenizer=tokenizer,
-        device="cpu"
+        model="IlyaGusev/saiga_llama3_8b",
+        device_map="auto"
     )
-    
-    return pipe
 
 # Поиск релевантной информации
 def find_relevant_info(question, collection, embed_fn, n_results=3):
@@ -106,12 +98,11 @@ def generate_response(question, context, llm_pipe):
     
     output = llm_pipe(
         prompt,
-        max_new_tokens=512,
+        max_new_tokens=256,
         do_sample=True,
         temperature=0.7,
         top_p=0.9,
-        repetition_penalty=1.2,
-        eos_token_id=2
+        repetition_penalty=1.2
     )
     
     return output[0]["generated_text"][len(prompt):].strip()