Update app.py

app.py

@@ -1,42 +1,43 @@
 import torch
 import gradio as gr
-from transformers import AutoTokenizer, AutoModelForSeq2SeqLM
+from transformers import AutoTokenizer, AutoModelForCausalLM, AutoConfig
 
-model_id = "google/flan-t5-base"  # можно flan-t5-large, если хватает памяти
-tokenizer = AutoTokenizer.from_pretrained(model_id)
-model = AutoModelForSeq2SeqLM.from_pretrained(model_id)
-device = "cuda" if torch.cuda.is_available() else "cpu"
-model.to(device)
+model_path = "cody82/unitrip"  # путь к локальной модели
 
-context = """
-Университет Иннополис был основан в 2012 году. Это современный вуз в России, 
-специализирующийся на IT и робототехнике, расположенный в городе Иннополис, Татарстан.
-"""
+config = AutoConfig.from_pretrained(model_path, local_files_only=True)
+tokenizer = AutoTokenizer.from_pretrained(model_path, local_files_only=True)
+model = AutoModelForCausalLM.from_pretrained(
+    model_path,
+    config=config,
+    local_files_only=True,
+    torch_dtype=torch.float32,
+    device_map="auto" if torch.cuda.is_available() else None,
+)
+
+system_message = "Ты — умный помощник по Университету Иннополис."
 
 def respond(message, history=None):
     if history is None:
         history = []
 
-    prompt = (
-        "Используя следующий контекст, ответь на вопрос четко и кратко.\n"
-        f"Контекст: {context}\n"
-        f"Вопрос: {message}\n"
-        "Ответ:"
-    )
-
-    inputs = tokenizer(prompt, return_tensors="pt").to(device)
+    prompt = f"{system_message}\nUser: {message}\nAssistant:"
 
+    inputs = tokenizer(prompt, return_tensors="pt").to(model.device)
     with torch.no_grad():
         outputs = model.generate(
             **inputs,
-            max_new_tokens=100,
+            max_new_tokens=128,
             do_sample=False,
-            eos_token_id=tokenizer.eos_token_id
+            pad_token_id=tokenizer.eos_token_id,
+            eos_token_id=tokenizer.eos_token_id,
+            use_cache=True,
         )
-    answer = tokenizer.decode(outputs[0], skip_special_tokens=True)
+
+    generated_tokens = outputs[0][inputs["input_ids"].shape[1]:]
+    answer = tokenizer.decode(generated_tokens, skip_special_tokens=True).strip()
 
     history.append((message, answer))
     return history
 
-iface = gr.ChatInterface(fn=respond, title="Innopolis Q&A")
-iface.launch()
+chat = gr.ChatInterface(fn=respond, title="Innopolis Assistant")
+chat.launch()