Update app.py

app.py

@@ -1,51 +1,30 @@
 import streamlit as st
 from transformers import AutoTokenizer, AutoModelForCausalLM
 import torch
-import random
+from PIL import Image
+import io
 
 @st.cache_resource
 def load_model():
     model_name = "Qwen/Qwen2-VL-7B-Instruct"
     tokenizer = AutoTokenizer.from_pretrained(model_name)
-    model = AutoModelForCausalLM.from_pretrained(model_name)
+    model = AutoModelForCausalLM.from_pretrained(model_name, device_map="auto")
     return tokenizer, model
 
-def generate_response(prompt, tokenizer, model):
-    inputs = tokenizer.encode(prompt, return_tensors='pt')
+def generate_response(prompt, image, tokenizer, model):
+    if image:
+        image = Image.open(image).convert('RGB')
+        inputs = tokenizer.from_pretrained(prompt, images=[image], return_tensors='pt').to(model.device)
+    else:
+        inputs = tokenizer(prompt, return_tensors='pt').to(model.device)
+    
     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 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 misspell_word(word):
-    if len(word) < 3:
-        return word
-    vowels = 'аеёиоуыэюя'
-    consonants = 'бвгджзйклмнпрстфхцчшщ'
+        outputs = model.generate(**inputs, max_new_tokens=100)
     
-    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
+    response = tokenizer.decode(outputs[0], skip_special_tokens=True)
+    return response
 
-st.title("AI Чат с простой русской моделью")
+st.title("Чат с Qwen VL-7B-Instruct")
 
 tokenizer, model = load_model()
 
@@ -55,14 +34,27 @@ if "messages" not in st.session_state:
 for message in st.session_state.messages:
     with st.chat_message(message["role"]):
         st.markdown(message["content"])
-
-if prompt := st.chat_input("Введите ваше сообщение"):
-    st.session_state.messages.append({"role": "user", "content": prompt})
-    with st.chat_message("user"):
-        st.markdown(prompt)
+        if "image" in message:
+            st.image(message["image"])
+
+prompt = st.chat_input("Введите ваше сообщение")
+uploaded_file = st.file_uploader("Загрузите изображение (необязательно)", type=["png", "jpg", "jpeg"])
+
+if prompt or uploaded_file:
+    if uploaded_file:
+        image = Image.open(uploaded_file)
+        st.session_state.messages.append({"role": "user", "content": prompt or "Опишите это изображение", "image": uploaded_file})
+        with st.chat_message("user"):
+            if prompt:
+                st.markdown(prompt)
+            st.image(image)
+    else:
+        st.session_state.messages.append({"role": "user", "content": prompt})
+        with st.chat_message("user"):
+            st.markdown(prompt)
     
     with st.chat_message("assistant"):
-        response = generate_response(prompt, tokenizer, model)
+        response = generate_response(prompt, uploaded_file, tokenizer, model)
         st.markdown(response)
     
     st.session_state.messages.append({"role": "assistant", "content": response})