File size: 2,576 Bytes
cb2c1f7
fcb3be6
ef23a53
d3eb12f
ef23a53
 
75c478f
cb2c1f7
603c09c
d5f380b
18586c6
 
 
5284e0b
18586c6
5284e0b
18586c6
d5f380b
18586c6
d5f380b
 
5284e0b
18586c6
 
d5f380b
ef23a53
cb2c1f7
 
ef23a53
d5f380b
 
 
cb2c1f7
d5f380b
ef23a53
d5f380b
 
 
 
 
 
 
5284e0b
 
d5f380b
 
fcb3be6
18586c6
d5f380b
 
fcb3be6
d5f380b
 
18586c6
5284e0b
d5f380b
fcb3be6
d5f380b
 
18586c6
 
 
d5f380b
 
 
18586c6
d5f380b
18586c6
cb2c1f7
18586c6
cb2c1f7
 
ef23a53
 
d5f380b
 
cb2c1f7
5284e0b
 
d5f380b
ef23a53
cb2c1f7
18586c6
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
from flask import Flask, request, jsonify
from transformers import pipeline
import logging

logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)

app = Flask(__name__)

# Инициализация pipeline
try:
    chatbot = pipeline(
        "text-generation",
        model="sberbank-ai/rugpt3medium_based_on_gpt2",
        device_map="auto",
        model_kwargs={"torch_dtype": "auto"}
    )
    logger.info("Model loaded successfully")
except Exception as e:
    logger.error(f"Error: {e}")
    chatbot = None

@app.route('/health', methods=['GET'])
def health_check():
    return jsonify({"status": "healthy", "model_loaded": chatbot is not None})

@app.route('/chat', methods=['POST'])
def chat():
    try:
        if not chatbot:
            return jsonify({"error": "Model not loaded"}), 500
            
        data = request.get_json()
        user_message = data.get('message', '').strip()
        
        if not user_message:
            return jsonify({"error": "Empty message"}), 400

        logger.info(f"Received: {user_message}")

        # Четкий промпт для адекватного общения
        prompt = f"""Ты - профессиональный помощник по Telegram стикерам. Отвечай вежливо и по делу.

Пользователь: {user_message}
Помощник:"""

        response = chatbot(
            prompt,
            max_new_tokens=100,
            temperature=0.7,
            do_sample=True,
            top_p=0.9,
            repetition_penalty=1.1,
            num_return_sequences=1,
            pad_token_id=chatbot.tokenizer.eos_token_id,
            truncation=True
        )

        # Извлекаем ответ
        generated_text = response[0]['generated_text']
        assistant_response = generated_text.replace(prompt, "").strip()
        
        # Очищаем ответ
        if "Пользователь:" in assistant_response:
            assistant_response = assistant_response.split("Пользователь:")[0].strip()
        
        logger.info(f"Response: {assistant_response}")
        
        return jsonify({
            "response": assistant_response,
            "status": "success"
        })
        
    except Exception as e:
        logger.error(f"Error: {e}")
        return jsonify({"error": "Internal error"}), 500

@app.route('/', methods=['GET'])
def home():
    return jsonify({"message": "Sticker assistant API"})

if __name__ == '__main__':
    app.run(host='0.0.0.0', port=7860, debug=False)