Update app.py

app.py

@@ -1,11 +1,10 @@
-import torch
 import gradio as gr
-from transformers import AutoTokenizer, AutoModelForSeq2SeqLM
-import re
+from transformers import AutoTokenizer, AutoModelForQuestionAnswering
+import torch
 
-model_id = "google/flan-t5-base"
-tokenizer = AutoTokenizer.from_pretrained(model_id)
-model = AutoModelForSeq2SeqLM.from_pretrained(model_id)
+model_name = "deepset/roberta-base-squad2"
+tokenizer = AutoTokenizer.from_pretrained(model_name)
+model = AutoModelForQuestionAnswering.from_pretrained(model_name)
 device = "cuda" if torch.cuda.is_available() else "cpu"
 model.to(device)
 
@@ -15,39 +14,22 @@ context = (
     "расположенный в городе Иннополис, Татарстан."
 )
 
-def clean_answer(answer, prompt):
-    # Убираем prompt из начала, если остался
-    answer = answer[len(prompt):].strip() if answer.lower().startswith(prompt.lower()) else answer.strip()
-    # Оставляем только кириллицу, пробелы и знаки препинания
-    answer = re.sub(r"[^а-яА-ЯёЁ ,.\-:;?!]", "", answer)
-    # Дополнительно можно убрать повторяющиеся символы
-    answer = re.sub(r"(.)\1{2,}", r"\1", answer)
-    return answer
-
-def respond(message, history=None):
+def respond(question, history=None):
     if history is None:
         history = []
 
-    prompt = (
-        "Используя следующий контекст, ответь на вопрос четко и кратко.\n"
-        f"Контекст: {context}\n"
-        f"Вопрос: {message}\n"
-        "Ответ:"
-    )
-
-    inputs = tokenizer(prompt, return_tensors="pt").to(device)
-
+    inputs = tokenizer.encode_plus(question, context, return_tensors="pt").to(device)
     with torch.no_grad():
-        outputs = model.generate(
-            **inputs,
-            max_new_tokens=100,
-            do_sample=False,
-            eos_token_id=tokenizer.eos_token_id
-        )
-    raw_answer = tokenizer.decode(outputs[0], skip_special_tokens=True)
-    answer = clean_answer(raw_answer, prompt)
+        outputs = model(**inputs)
+    start_scores = outputs.start_logits
+    end_scores = outputs.end_logits
+
+    start = torch.argmax(start_scores)
+    end = torch.argmax(end_scores) + 1
+    answer_tokens = inputs['input_ids'][0][start:end]
+    answer = tokenizer.decode(answer_tokens, skip_special_tokens=True)
 
-    history.append((message, answer))
+    history.append((question, answer))
     return history
 
 iface = gr.ChatInterface(fn=respond, title="Innopolis Q&A")