gradio

app.py

@@ -1,7 +1,7 @@
 import transformers
 import torch
 import tokenizers
-import streamlit as st
+import gradio as gr
 import re
 
 from PIL import Image
@@ -20,7 +20,7 @@ def get_model_t5(model_name,tokenizer_name):
     return model, tokenizer
 
 
-def predict_gpt(text, model, tokenizer, temperature):
+def predict_gpt(text, model, tokenizer, temperature=1.0):
     input_ids = tokenizer.encode(text, return_tensors="pt")
     with torch.no_grad():
         out = model.generate(input_ids, 
@@ -40,7 +40,7 @@ def predict_gpt(text, model, tokenizer, temperature):
     generated_text = list(map(decode, out['sequences'])).split('Описание :')[1]
     return generated_text
 
-def predict_t5(text, model, tokenizer, temperature):
+def predict_t5(text, model, tokenizer, temperature=1.2):
     input_ids = tokenizer.encode(text, return_tensors="pt")
     with torch.no_grad():
         out = model.generate(input_ids, 
@@ -59,30 +59,27 @@ def predict_t5(text, model, tokenizer, temperature):
     generated_text = list(map(decode, out['sequences']))[0]
     return generated_text
 
+def generate(model,text):
+    if model == 'GPT':
+                result = predict_gpt(text, gpt_model, gpt_tokenizer)
+    elif model == 'T5':
+                result = predict_t5(text, t5_model, t5_tokenizer)
+    return result
+         
+
 gpt_model, gpt_tokenizer = get_model_gpt('mipatov/rugpt3_nb_descr', 'mipatov/rugpt3_nb_descr')
 t5_model, t5_tokenizer = get_model_t5('mipatov/rut5_nb_descr', 'mipatov/rut5_nb_descr')
 
 
-option = st.selectbox('Выберите модель', ('GPT', 'T5'))
-temperature = st.slider(label='Температура', min_value=0.1, max_value=5.0, value=1.0,step = 0.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 ;'
+demo = gr.Interface(
+    fn=generate,
+    inputs=[
+        gr.components.Dropdown(label="Модель", choices=('GPT', 'T5')),
+        gr.components.Textbox(label="Характеристики ноутбука",value = example,lines = example.count('\n')),
+    ],
+    outputs=["text"],
+)
 
-text = st.text_area(label='Характеристики ноутбука', value=example, height=200).replace('\n','')
-button = st.button('Старт')
+demo.launch()
 
-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: @mipatov")