# تثبيت المكتبات المطلوبة
!pip install gradio transformers accelerate bitsandbytes huggingface_hub

import gradio as gr
import torch
from transformers import AutoTokenizer, AutoModelForCausalLM, BitsAndBytesConfig
from huggingface_hub import login


# 🔹 تحديد اسم النموذج
model_name = "google/gemma-2b-it"

# 🔹 تحميل التوكن
tokenizer = AutoTokenizer.from_pretrained(model_name)

# 🔹 تحميل النموذج مع تحسين استهلاك الذاكرة
quantization_config = BitsAndBytesConfig(load_in_8bit=True)  # تحميل النموذج بتنسيق 8-bit لتقليل استهلاك الذاكرة

model = AutoModelForCausalLM.from_pretrained(
    model_name,
    quantization_config=quantization_config,
    device_map="auto"  # توزيع الحمل بين GPU و RAM تلقائيًا
)



print(f"✅ تم رفع النموذج بنجاح! يمكنك الوصول إليه على: https://huggingface.co/{username}/{repo_name}")

# 🔹 دالة معالجة الأسئلة الخاصة بالسباحة
def swimming_coach_advice(question):
    prompt = f"User: {question}\nAI Coach:"

    # تحويل النص إلى أرقام
    inputs = tokenizer(prompt, return_tensors="pt").to("cuda")  # استخدام الـ GPU إذا كان متاحًا

    # توليد النص بطريقة فعالة
    output = model.generate(
        **inputs,
        max_length=150,  # تحديد الطول الأقصى للخروج
        temperature=0.7,  # ضبط العشوائية
        top_p=0.9,  # تحسين الاستجابة
        repetition_penalty=1.2  # تقليل التكرار
    )

    # تحويل المخرجات إلى نص
    response = tokenizer.decode(output[0], skip_special_tokens=True)

    return response

# 🔹 إنشاء واجهة Gradio
interface = gr.Interface(
    fn=swimming_coach_advice,
    inputs=gr.Textbox(label="Ask the AI Swimming Coach a question"),
    outputs=gr.Textbox(label="AI Coach Response"),
    title="AI Swimming Coach",
    description="Ask me anything about swimming techniques, training, or tips!"
)

# 🔹 تشغيل التطبيق
if __name__ == "__main__":
    interface.launch()