from transformers import AutoConfig, AutoModelForCausalLM, AutoTokenizer, pipeline
import gradio as gr

model_path = "interneuronai/az-llama2"
model = AutoModelForCausalLM.from_pretrained(model_path)
tokenizer = AutoTokenizer.from_pretrained(model_path)
pipe = pipeline(task="text-generation", model=model, tokenizer=tokenizer, max_length=300)

def get_response(user_input):
    instruction = f"Sən psixoloq rolundasan. İstifadəçi deyir: '{user_input}'. Empatik və dəstəkləyici bir cavab ver, amma professional psixoloji yardımın yerini tuta bilməyəcəyini vurğula."
    formatted_prompt = f"""Aşağıda daha çox kontekst təmin edən təlimat var. Sorğunu adekvat şəkildə tamamlayan cavab yazın.
                    ### Təlimat:
                    {instruction}
                    ### Cavab:
                    """
    result = pipe(formatted_prompt)
    response = result[0]['generated_text'].split("### Cavab:")[-1].strip()
    return response

def chat(message, history):
    response = get_response(message)
    return response

demo = gr.ChatInterface(
    chat,
    title="Psixoloji Məsləhət Assistenti",
    description="Bu bir AI asistentidir və professional psixoloji yardımın yerini tuta bilməz. Ciddi problemlər üçün həmişə mütəxəssisə müraciət edin.",
    theme=gr.themes.Soft(),
    examples=["Stress hiss edirəm", "Yuxu problemlərim var", "Özümü tənha hiss edirəm"],
    retry_btn="Yenidən cəhd et",
    undo_btn="Geri al",
    clear_btn="Təmizlə"
)

demo.launch()