nb app

app.py

@@ -1,7 +1,102 @@
-import gradio as gr
+import transformers
+import torch
+import tokenizers
+import streamlit as st
+import re
 
-def greet(name):
-    return "Hello " + name + "!!"
+from PIL import Image
 
-iface = gr.Interface(fn=greet, inputs="text", outputs="text")
-iface.launch()
+
+@st.cache(hash_funcs={tokenizers.Tokenizer: lambda _: None, tokenizers.AddedToken: lambda _: None, re.Pattern: lambda _: None}, allow_output_mutation=True, suppress_st_warning=True)
+def get_model(model_name, model_path):
+    tokenizer = transformers.GPT2Tokenizer.from_pretrained(model_name)
+    model = transformers.GPT2LMHeadModel.from_pretrained(model_name)
+    model.eval()
+    return model, tokenizer
+
+
+def predict(text, model, tokenizer, n_beams=5, temperature=2.5, top_p=0.8, length_of_generated=300):
+    # text += '\n'
+    input_ids = tokenizer.encode(text, return_tensors="pt")
+    length_of_prompt = len(input_ids[0])
+    with torch.no_grad():
+        out = model.generate(input_ids,
+                             do_sample=True,
+                             num_beams=n_beams,
+                             temperature=temperature,
+                             top_p=top_p,
+                             max_length=length_of_prompt + length_of_generated,
+                             eos_token_id=tokenizer.eos_token_id
+                             )
+                          
+    generated = list(map(tokenizer.decode, out))[0]
+    return generated.replace('\n[EOS]\n', '')
+
+def predict_gpt(text, model, tokenizer,):
+    input_ids = tokenizer.encode(text, return_tensors="pt").to(device)
+    with torch.no_grad():
+        out = model.generate(input_ids, 
+                            do_sample=True,
+                            num_beams=3,
+                            temperature=1.0,
+                            top_p=0.75,
+                            max_length=1024,
+                            eos_token_id = tokenizer.eos_token_id,
+                            pad_token_id = tokenizer.pad_token_id,
+                            repetition_penalty = 2.5,
+                            num_return_sequences = 1,
+                            output_attentions = True,
+                            return_dict_in_generate=True,
+                            )
+    decode = lambda x : tokenizer.decode(x, skip_special_tokens=True)
+    generated_text = list(map(decode, out['sequences'])).split('Описание :')[1]
+    return generated_text
+
+def predict_t5(text, model, tokenizer,):
+    input_ids = tokenizer.encode(text, return_tensors="pt").to(device)
+    with torch.no_grad():
+        out = model.generate(input_ids, 
+                        do_sample=True,
+                        num_beams=4,
+                        temperature=1.2,
+                        top_p=0.35,
+                        max_length=1024,
+                        length_penalty = 5.5,
+                        output_attentions = True,
+                        return_dict_in_generate=True,
+                        repetition_penalty = 2.5,
+                        num_return_sequences = 1
+                        )
+    decode = lambda x : tokenizer.decode(x, skip_special_tokens=True)
+    generated_text = list(map(decode, out['sequences']))[0]
+    return generated_text
+
+gpt_model, gpt_tokenizer = get_model('mipatov/rugpt3_nb_descr', 'mipatov/rugpt3_nb_descr')
+t5_model, t5_tokenizer = get_model('mipatov/rugpt3_nb_descr', 'mipatov/rugpt3_nb_descr')
+
+# st.title("NeuroKorzh")
+
+
+option = st.selectbox('Выберите модель', ('GPT', 'T5'))
+temperature = st.slider(label='Температура', min_value=0.1, max_value=10, value=1)
+
+# st.markdown("\n")
+example = '  Название : Super_NB 2001 Gaming;\n Диагональ экрана : 21 " ;\n Разрешение экрана : 1337x228 ;\n Поверхность экрана : матовая ;\n Тип матрицы : nfc ;\n Процессор : CMD processor 7 ядер 16.16 ГГц (46.0 ГГц, в режиме Turbo) ;\n Оперативная память : 28 Гб DDR5 ;\n Тип видеокарты : интегрированный ;\n Видеокарта : qwerty-grafics ;\n SSD : 720 Гб ;\n Wi-Fi : ДА, 802.11 a/b/g/n/ac ;\n Bluetooth : ДА, v5.0 ;\n Кабельная сеть : 10/100/1000 (Gigabit Ethernet) Мбит/с ;\n USB 2.0 : 13 ;\n USB 3.0 : 22 ;\n HDMI : 11 ;\n Операционная система : CMD-shell ;\n Веб-камера : встроенная ;\n Микрофон : есть ;\n Разъем наушники/микрофон : комбинированный разъем ;\n Акустическая система : стереодинамики ;\n Цвет клавиатуры : черный ;\n Цифровой блок клавиатуры : есть ;\n Подсветка клавиш клавиатуры : есть ;\n Тип батареи : Al-Ion ;\n Количество ячеек батареи : 36 cell ;\n Энергоемкость батареи : 176 Wh ;\n Цвет : черный ;\n Размеры : 1.23 х  456 х  78.9 мм ;\n Вес : 19 кг ;\n Гарантия : 322 мес. ;\n Материал корпуса : пластик ;\n Время работы от батареи : 82ч ;\n Кард-ридер : есть WA SD ;'
+
+text = st.text_area(label='Характеристики ноутбука', value=example, height=200).replace('\n','')
+button = st.button('Старт')
+
+if button:
+    try:
+        with st.spinner("Пишем описание..."):
+            if option == 'GPT':
+                result = predict_gpt(text, gpt_model, gpt_tokenizer, temperature=temperature)
+            elif option == 'T5':
+                result = predict_t5(text, t5_model, t5_tokenizer, temperature=temperature)
+            else:
+                st.error('Error in selectbox')
+        
+        st.text_area(label='', value=result, height=1000)
+    
+    except Exception:
+        st.error("Ooooops, something went wrong. Please try again and report to me, tg: @vladyur")