Update app.py

app.py

@@ -8,8 +8,9 @@ from torch import tensor
 
 import joblib
 from dataclasses import dataclass
-from transformers import AutoTokenizer, AutoModel, AutoModelForSequenceClassification
+from transformers import AutoTokenizer, AutoModel, AutoModelForSequenceClassification, GPT2LMHeadModel, GPT2Tokenizer
 import json
+import os
 
 from preprocessing import predict_review, data_preprocessing_hard
 from model_lstm import LSTMClassifier
@@ -173,17 +174,39 @@ elif selected_model == "Оценка степени токсичности по
     
 
 
-# Генерация текста GPT-моделью
+# Генерация текста GPT
 elif selected_model == "Генерация текста GPT-моделью по пользовательскому prompt":
     st.title("""
-    Приложение генерирует текст по Вашему промту
+    Нейросетевой гороскоп Якубовский-Дьяченко
     """)
 
     st.write("""
-    Для генерации текста используется предобученная сеть GPT.
+    Для генерации текста используется предобученная сеть GPT2. Дообучение проходило на гороскопах. 
+    Общая длина текста для обучения 37 001 887 слов.
     """)
-    uploaded_img = st.sidebar.file_uploader('Загрузи свое космофото', type=["jpg", "png", "jpeg"])
-    if uploaded_img is not None:
-        input_img = io.imread(uploaded_img)
-    else:
-        input_img = io.imread('/Users/id/Documents/strlit/cv_project/Segm.jpg')
+    user_text_input = st.text_area('Введите информацию о себе для формиорования гороскопа:')
+
+    # GPT2
+    model_path = "model.safetensors"
+    huggingface_token = os.getenv("HF_TOKEN")
+    model = GPT2LMHeadModel.from_pretrained(model_path, token=huggingface_token)
+    tokenizer = GPT2Tokenizer.from_pretrained(model_path, token=huggingface_token)
+
+    if st.button('Сделать гороскоп'):
+        start_time = time.time()
+        input_ids = tokenizer.encode(user_text_input, return_tensors="pt").to(device)
+        model.eval()
+        with torch.no_grad():
+            out = model.generate(input_ids, 
+                                do_sample=True,
+                                num_beams=2,
+                                temperature=1.1,
+                                top_p=0.9,
+                                max_length=50,
+                                )
+
+        generated_text = list(map(tokenizer.decode, out))[0]
+        end_time = time.time()
+        prediction_time = end_time - start_time
+
+        st.write(f'Ваше предсказание: {generated_text}')