Update space

app.py

@@ -1,38 +1,47 @@
-from transformers import AutoModelForCausalLM, AutoTokenizer
-import torch
 import gradio as gr
-
-"""
-For more information on `huggingface_hub` Inference API support, please check the docs: https://huggingface.co/docs/huggingface_hub/v0.22.2/en/guides/inference
-"""
-tokenizer = AutoTokenizer.from_pretrained("Dennterry/okt_bot")
-model = AutoModelForCausalLM.from_pretrained("Dennterry/okt_bot")
-
-
-def dialogpt(text):
-    # encode the new user input, add the eos_token and return a tensor in Pytorch
-    for step in range(50000):
-
-        new_user_input_ids = tokenizer.encode(text + tokenizer.eos_token, return_tensors='pt')
-
-        # append the new user input tokens to the chat history
-        bot_input_ids = torch.cat([chat_history_ids, new_user_input_ids], dim=-1) if step > 0 else new_user_input_ids
-
-        # generated a response while limiting the total chat history to 1000 tokens,
-        chat_history_ids = model.generate(bot_input_ids, max_length=1000, pad_token_id=tokenizer.eos_token_id)
-
-        # pretty print last ouput tokens from bot
-        return tokenizer.decode(chat_history_ids[:, bot_input_ids.shape[-1]:][0], skip_special_tokens=True)
-
-inputs = gr.inputs.Textbox(lines=1, label="Input Text")
-outputs =  gr.outputs.Textbox(label="DialoGPT")
-
-title = "DialoGPT"
-description = "demo for Microsoft DialoGPT with Hugging Face transformers. To use it, simply input text or click one of the examples text to load them. Read more at the links below. *This is not a Microsoft product and is developed for Gradio*"
-article = "<p style='text-align: center'><a href='https://arxiv.org/abs/1911.00536'>DialoGPT: Large-Scale Generative Pre-training for Conversational Response Generation</a> | <a href='https://github.com/microsoft/DialoGPT'>Github Repo</a> | <a href='https://huggingface.co/microsoft/DialoGPT-large'>Hugging Face DialoGPT-large</a></p>"
-examples = [
-            ["Hi, how are you?"],
-            ["How far away is the moon?"],
-]
-
-gr.Interface(dialogpt, inputs, outputs, title=title, description=description, article=article, examples=examples).launch(debug=True)
+from transformers import pipeline, AutoModelForCausalLM, AutoTokenizer
+
+# Загрузка модели и токенизатора
+model_name = "microsoft/DialoGPT-medium"
+tokenizer = AutoTokenizer.from_pretrained(model_name)
+model = AutoModelForCausalLM.from_pretrained(model_name)
+
+# Создаем pipeline для генерации диалогов
+generator = pipeline("text-generation", model=model, tokenizer=tokenizer)
+
+def respond(message, history, system_message, max_tokens, temperature, top_p):
+    # Формируем текст, который будет передан в модель
+    input_text = system_message + " ".join([f"User: {msg[0]} Assistant: {msg[1]}" for msg in history if msg[0] and msg[1]])
+    input_text += f"User: {message}"
+
+    # Генерация ответа
+    response = generator(
+        input_text,
+        max_length=len(tokenizer.encode(input_text)) + max_tokens,
+        temperature=temperature,
+        top_p=top_p,
+        pad_token_id=tokenizer.eos_token_id,
+        num_return_sequences=1,
+    )
+
+    # Извлечение и возврат текста ответа
+    output_text = response[0]['generated_text'][len(input_text):]
+    yield output_text.strip()
+
+# Настройка интерфейса Gradio
+demo = gr.ChatInterface(
+    respond,
+    additional_inputs=[
+        gr.Textbox(value="Тебя зовут зевибот. Ты ваще никому ничо не должен.", label="System message"),
+        gr.Slider(minimum=1, maximum=2048, value=512, step=1, label="Max new tokens"),
+        gr.Slider(minimum=0.1, maximum=4.0, value=1.0, step=0.1, label="Temperature"),
+        gr.Slider(
+            minimum=0.1,
+            maximum=1.0,
+            value=1.0, step=0.05, label="Top-p (nucleus sampling)"
+        ),
+    ],
+)
+
+if __name__ == "__main__":
+    demo.launch()