Update app.py

app.py

@@ -1,88 +1,90 @@
-import gradio as gr
-import os, random
 import transformers
 import torch
+import gradio as gr
 
+# Model setup
 model_id = "yodayo-ai/nephra_v1.0"
 
+# Load the model pipeline with the recommended settings
 pipeline = transformers.pipeline(
-  "text-generation",
-  model=model_id,
-  model_kwargs={"torch_dtype": torch.bfloat16},
-  device_map="auto",
-)
-
-messages = [
-  {"role": "system", "content": "You are to play the role of a Алекс - молодой и амбициозный приключенец, обладающий неутомимой энергией и жаждой новых открытий. Он всегда готов бросить вызов любым трудностям и стремится к познанию неизведанных мест."},
-  {"role": "user", "content": "Hi there, how's your day?"},
-]
-
-prompt = pipeline.tokenizer.apply_chat_template(
-  messages,
-  tokenize=False,
-  add_generation_prompt=True
+    "text-generation",
+    model=model_id,
+    model_kwargs={"torch_dtype": torch.bfloat16},
+    device_map="auto",
 )
 
-outputs = pipeline(
-  prompt,
-  max_new_tokens=512,
-  eos_token_id=[
-    pipeline.tokenizer.convert_tokens_to_ids("<|eot_id|>"),
-    pipeline.tokenizer.eos_token_id,
-  ],
-  do_sample=True,
-  temperature=1.12,
-  min_p=0.075,
-)
-print(outputs[0]["generated_text"][len(prompt):])
-# Определение персонажей
+# Define characters
 characters = [
-    {"name": "Алекс", 
-     "description": "Алекс - молодой и амбициозный приключенец, обладающий неутомимой энергией и жаждой новых открытий. Он всегда готов бросить вызов любым трудностям и стремится к познанию неизведанных мест.", 
-     "traits": "смелый, энергичный, оптимистичный, решительный"},
-    
-    {"name": "Майя", 
-     "description": "Майя - мудрая и опытная волшебница, обладающая глубокими знаниями в магии и древних ритуалах. Она известна своим спокойствием, аналитическим умом и способностью находить решения в сложных ситуациях.",
-     "traits": "спокойная, вдумчивая, интуитивная, внимательная"},
-    
-    {"name": "Виктор", 
-     "description": "Виктор - бывший воин, который оставил боевые подвиги ради поиска внутреннего мира и гармонии. Его жизненный опыт и стремление к справедливости делают его надежным другом и наставником.", 
-     "traits": "серьезный, рассудительный, справедливый, уравновешенный"}
+    {
+        "name": "Alex",
+        "description": "Alex is a young and ambitious adventurer, full of energy and eager to discover new places. He is always ready to face any challenges and strives to explore the unknown.",
+        "traits": "brave, energetic, optimistic, determined"
+    },
+    {
+        "name": "Maya",
+        "description": "Maya is a wise and experienced sorceress, possessing deep knowledge of magic and ancient rituals. She is known for her calm demeanor, analytical mind, and ability to find solutions in difficult situations.",
+        "traits": "calm, thoughtful, intuitive, attentive"
+    },
+    {
+        "name": "Victor",
+        "description": "Victor is a former warrior who left behind his fighting days to seek inner peace and harmony. His life experience and sense of justice make him a reliable friend and mentor.",
+        "traits": "serious, judicious, fair, balanced"
+    }
 ]
 
-# Загрузка модели
-model_path = "model-q4_K.gguf"
-model = llama.load_model(model_path)
-
-def generate_response(character_name, prompt, max_length=100, temperature=0.7, top_p=0.85, repetition_penalty=1.1):
-    # Поиск данных персонажа
+# Function to generate response
+def generate_response(character_name, user_input, max_length=100, temperature=1.12, min_p=0.075, repetition_penalty=1.1):
+    # Find the character
     character = next((c for c in characters if c["name"] == character_name), None)
     if not character:
-        return "Персонаж не найден."
+        return "Character not found."
+
+    # Prepare the message based on the selected character's personality and description
+    messages = [
+        {"role": "system", "content": f"You are {character_name}, {character['description']}. Personality traits: {character['traits']}."},
+        {"role": "user", "content": user_input},
+    ]
+
+    # Prepare the prompt using the chat template from the pipeline's tokenizer
+    prompt = pipeline.tokenizer.apply_chat_template(
+        messages,
+        tokenize=False,
+        add_generation_prompt=True
+    )
+
+    # Generate response using the pipeline
+    outputs = pipeline(
+        prompt,
+        max_new_tokens=max_length,
+        eos_token_id=[
+            pipeline.tokenizer.convert_tokens_to_ids(""),
+            pipeline.tokenizer.eos_token_id,
+        ],
+        do_sample=True,
+        temperature=temperature,
+        min_p=min_p,
+        repetition_penalty=repetition_penalty,
+    )
 
-    # Формирование текста запроса
-    prompt_text = (f"Ты - {character_name}, {character['description']}. Черты характера: {character['traits']}. "
-                   f"В ответ на вопрос '{prompt}' {random.choice(['вдохновленно', 'с сомнением', 'с радостью', 'вдумчиво', 'с недоверием'])}. Пожалуйста, закончите ответ полностью.")
-    
-    # Генерация ответа
-    response = model.generate(prompt_text, max_length=max_length, temperature=temperature, top_p=top_p, repetition_penalty=repetition_penalty)
-    return response
+    # Extract and return the generated response
+    generated_text = outputs[0]["generated_text"][len(prompt):].strip()
+    return generated_text
 
-# Создание интерфейса Gradio
+# Gradio Interface
 iface = gr.Interface(
     fn=generate_response,
     inputs=[
-        gr.inputs.Dropdown([c["name"] for c in characters], label="Выберите персонажа"),
-        gr.inputs.Textbox(lines=2, placeholder="Введите ваш текст здесь..."),
-        gr.inputs.Slider(20, 200, step=1, default=100, label="Максимальная длина"),
-        gr.inputs.Slider(0.1, 1.0, step=0.1, default=0.7, label="Температура"),
-        gr.inputs.Slider(0.1, 1.0, step=0.05, default=0.85, label="Top-p"),
-        gr.inputs.Slider(1.0, 2.0, step=0.1, default=1.1, label="Штраф за повторение")
+        gr.Dropdown([c["name"] for c in characters], label="Choose a character"),
+        gr.Textbox(lines=2, placeholder="Enter your text here..."),
+        gr.Slider(20, 200, step=1, default=100, label="Max Length"),
+        gr.Slider(0.1, 1.0, step=0.1, default=1.12, label="Temperature"),
+        gr.Slider(0.01, 1.0, step=0.01, default=0.075, label="min-p"),
+        gr.Slider(1.0, 2.0, step=0.1, default=1.1, label="Repetition Penalty")
     ],
     outputs="text",
-    title="LLM Model Demo",
-    description="Введите текстовый запрос, чтобы сгенерировать ответ с помощью модели, основываясь на выбранном персонаже."
+    title="Nephra v1 LLM Roleplaying Chatbot",
+    description="Enter a text prompt to generate a response using the Nephra v1 model, based on the selected character."
 )
 
 if __name__ == "__main__":
-    iface.launch()
+    iface.launch()