Переход на Sage

app.py

@@ -3,8 +3,7 @@ import time
 import json
 import requests
 import gradio as gr
-import torch
-from transformers import GPT2LMHeadModel, GPT2Tokenizer
+import poe
 
 with open("config.json", "r") as f:
     config = json.load(f)
@@ -15,41 +14,21 @@ max_attempts = config["MAX_ATTEMPS"]
 wait_time = config["WAIT_TIME"]
 chatgpt_url = config["CHATGPT_URL"]
 system_prompt = config["SYSTEM_PROMPT"]
-sber_gpt = config["SBER_GRT"]
-use_sber = config["USE_SBER"]
-
-DEVICE = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
-
-if use_sber:
-    tokenizer = GPT2Tokenizer.from_pretrained(sber_gpt)
-    model = GPT2LMHeadModel.from_pretrained(sber_gpt).to(DEVICE)
-
-def generate(
-    model, tok, text,
-    do_sample=True, max_length=10000, repetition_penalty=5.0,
-    top_k=5, top_p=0.95, temperature=1,
-    num_beams=10,
-    no_repeat_ngram_size=3
-    ):
-    input_ids = tok.encode(text, return_tensors="pt").to(DEVICE)
-    out = model.generate(
-        input_ids.to(DEVICE),
-        max_length=max_length,
-        repetition_penalty=repetition_penalty,
-        do_sample=do_sample,
-        top_k=top_k, top_p=top_p, temperature=temperature,
-        num_beams=num_beams, no_repeat_ngram_size=no_repeat_ngram_size
-        )
-    return list(map(tok.decode, out))[0]
+use_sage = config["USE_SAGE"]
+sage_token = config["SAGE_TOKEN"]
+
+client = poe.Client(sage_token)
 
 def get_answer(question: str) -> Dict[str, Any]:
-    if use_sber:
-        content = generate(model, tokenizer, question)
+
+    if use_sage:
+        for chunk in client.send_message("capybara", question, with_chat_break=True):
+            pass
         return {
             'status': True,
-            'content': content
+            'content': chunk["text"]
         }
-
+    
     headers = {
         'Content-Type': 'application/json; charset=utf-8'
     }

config.json

@@ -4,7 +4,7 @@
     "MAX_ATTEMPS": 5,
     "WAIT_TIME": 1,
     "CHATGPT_URL": "https://free.churchless.tech/v1/chat/completions",
-    "SBER_GRT": "ai-forever/rugpt3small_based_on_gpt2",
-    "USE_SBER": 0,
+    "USE_SAGE": 1,
+    "SAGE_TOKEN": "PGUXiyEZKRHcMoij9AjxXw%3D%3D",
     "SYSTEM_PROMPT": "Your task is to give the most detailed answer to the question posed. At the beginning of the question, there are tags in square brackets specifying the subject of the question. It is necessary to answer in the language of the user's question"
 }

requirements.txt

@@ -1,2 +1,3 @@
 torch
-transformers
+transformers
+poe-api