import logging
from telegram import Update
from telegram.ext import ApplicationBuilder, CommandHandler, MessageHandler, filters
from transformers import AutoTokenizer, AutoModelForCausalLM
import torch

# تنظیمات لاگ
logging.basicConfig(
    format='%(asctime)s - %(name)s - %(levelname)s - %(message)s',
    level=logging.INFO
)

# بارگیری مدل و توکنایزر
model_name = "نام مدل شما"  # نام مدل خود را اینجا وارد کنید
tokenizer = AutoTokenizer.from_pretrained(model_name)
model = AutoModelForCausalLM.from_pretrained(model_name)

# تابع برای تولید پاسخ
def generate_response(input_text):
    inputs = tokenizer(input_text, return_tensors="pt").to("cuda")
    output = model.generate(**inputs, max_new_tokens=100)
    response = tokenizer.decode(output[0], skip_special_tokens=True)
    return response

# تابع شروع
async def start(update: Update, context):
    await update.message.reply_text('سلام! من یک بات هوش مصنوعی هستم. می‌توانید با من چت کنید.')

# تابع پاسخ به پیام‌ها
async def echo(update: Update, context):
    user_message = update.message.text
    response = generate_response(user_message)
    await update.message.reply_text(response)

# تابع اصلی
def main():
    # توکن بات تلگرام
    token = "توکن بات شما"  # توکن بات خود را اینجا وارد کنید

    # ایجاد برنامه بات
    application = ApplicationBuilder().token(token).build()

    # اضافه کردن دستورات
    start_handler = CommandHandler('start', start)
    application.add_handler(start_handler)

    # اضافه کردن handler برای پیام‌های معمولی
    echo_handler = MessageHandler(filters.TEXT & ~filters.COMMAND, echo)
    application.add_handler(echo_handler)

    # شروع بات
    application.run_polling()

if __name__ == '__main__':
    main()