Update app.py

app.py

@@ -1,105 +1,68 @@
-import os
-os.system('pip install dashscope')
-import gradio as gr
-from http import HTTPStatus
-import dashscope
-from dashscope import Generation
-from dashscope.api_entities.dashscope_response import Role
-from typing import List, Optional, Tuple, Dict
-from urllib.error import HTTPError
-default_system = 'You are a helpful assistant.'
+import streamlit as st
+from transformers import AutoTokenizer, AutoModelForCausalLM
+import torch
+import random
 
-YOUR_API_TOKEN = os.getenv('YOUR_API_TOKEN')
-dashscope.api_key = YOUR_API_TOKEN
+@st.cache_resource
+def load_model():
+    model_name = "sberbank-ai/rugpt3small_based_on_gpt2"
+    tokenizer = AutoTokenizer.from_pretrained(model_name)
+    model = AutoModelForCausalLM.from_pretrained(model_name)
+    return tokenizer, model
 
-History = List[Tuple[str, str]]
-Messages = List[Dict[str, str]]
+def generate_response(prompt, tokenizer, model):
+    inputs = tokenizer.encode(prompt, return_tensors='pt')
+    with torch.no_grad():
+        outputs = model.generate(inputs, max_length=100, num_return_sequences=1, 
+                                 temperature=0.9, top_k=50, top_p=0.95)
+    response = tokenizer.decode(outputs[0], skip_special_tokens=True)
+    return add_mistakes(response)
 
-def clear_session() -> History:
-    return '', []
+def add_mistakes(text):
+    words = text.split()
+    for i in range(len(words)):
+        if random.random() < 0.2:  # 20% шанс ошибки в слове
+            words[i] = misspell_word(words[i])
+    return ' '.join(words)
 
-def modify_system_session(system: str) -> str:
-    if system is None or len(system) == 0:
-        system = default_system
-    return system, system, []
-
-def history_to_messages(history: History, system: str) -> Messages:
-    messages = [{'role': Role.SYSTEM, 'content': system}]
-    for h in history:
-        messages.append({'role': Role.USER, 'content': h[0]})
-        messages.append({'role': Role.ASSISTANT, 'content': h[1]})
-    return messages
-
-
-def messages_to_history(messages: Messages) -> Tuple[str, History]:
-    assert messages[0]['role'] == Role.SYSTEM
-    system = messages[0]['content']
-    history = []
-    for q, r in zip(messages[1::2], messages[2::2]):
-        history.append([q['content'], r['content']])
-    return system, history
-
-
-def model_chat(query: Optional[str], history: Optional[History], system: str
-) -> Tuple[str, str, History]:
-    if query is None:
-        query = ''
-    if history is None:
-        history = []
-    messages = history_to_messages(history, system)
-    messages.append({'role': Role.USER, 'content': query})
-    gen = Generation.call(
-        model='qwen2-72b-instruct',
-        messages=messages,
-        result_format='message',
-        stream=True
-    )
-    for response in gen:
-        if response.status_code == HTTPStatus.OK:
-            role = response.output.choices[0].message.role
-            response = response.output.choices[0].message.content
-            system, history = messages_to_history(messages + [{'role': role, 'content': response}])
-            yield '', history, system
-        else:
-            raise ValueError('Request id: %s, Status code: %s, error code: %s, error message: %s' % (
-                response.request_id, response.status_code,
-                response.code, response.message
-            ))
+def misspell_word(word):
+    if len(word) < 3:
+        return word
+    vowels = 'аеёиоуыэюя'
+    consonants = 'бвгджзйклмнпрстфхцчшщ'
+    
+    if random.random() < 0.5:
+        # Заменяем случайную гласную
+        for i, char in enumerate(word):
+            if char.lower() in vowels:
+                replacement = random.choice(vowels)
+                return word[:i] + replacement + word[i+1:]
+    else:
+        # Заменяем случайную согласную
+        for i, char in enumerate(word):
+            if char.lower() in consonants:
+                replacement = random.choice(consonants)
+                return word[:i] + replacement + word[i+1:]
+    return word
 
+st.title("AI Чат с простой русской моделью")
 
-with gr.Blocks() as demo:
-    gr.Markdown("""<center><font size=8>Qwen2-72B-instruct Chat👾</center>""")
+tokenizer, model = load_model()
 
-    with gr.Row():
-        with gr.Column(scale=3):
-            system_input = gr.Textbox(value=default_system, lines=1, label='System')
-        with gr.Column(scale=1):
-            modify_system = gr.Button("🛠️ Set system prompt and clear history", scale=2)
-        system_state = gr.Textbox(value=default_system, visible=False)
-    chatbot = gr.Chatbot(label='qwen2-72B-instruct')
-    textbox = gr.Textbox(lines=1, label='Input')
+if "messages" not in st.session_state:
+    st.session_state.messages = []
 
-    with gr.Row():
-        clear_history = gr.Button("🧹 Clear history")
-        sumbit = gr.Button("🚀 Send")
+for message in st.session_state.messages:
+    with st.chat_message(message["role"]):
+        st.markdown(message["content"])
 
-    textbox.submit(model_chat,
-                 inputs=[textbox, chatbot, system_state],
-                 outputs=[textbox, chatbot, system_input],
-                 concurrency_limit = 40)
+if prompt := st.chat_input("Введите ваше сообщение"):
+    st.session_state.messages.append({"role": "user", "content": prompt})
+    with st.chat_message("user"):
+        st.markdown(prompt)
     
-    sumbit.click(model_chat,
-                 inputs=[textbox, chatbot, system_state],
-                 outputs=[textbox, chatbot, system_input],
-                 concurrency_limit = 40)
-    clear_history.click(fn=clear_session,
-                        inputs=[],
-                        outputs=[textbox, chatbot],
-                        concurrency_limit = 40)
-    modify_system.click(fn=modify_system_session,
-                        inputs=[system_input],
-                        outputs=[system_state, system_input, chatbot],
-                        concurrency_limit = 40)
-
-demo.queue(api_open=False)
-demo.launch(max_threads=40)
+    with st.chat_message("assistant"):
+        response = generate_response(prompt, tokenizer, model)
+        st.markdown(response)
+    
+    st.session_state.messages.append({"role": "assistant", "content": response})