import gradio as gr
import torch
from transformers import AutoTokenizer, AutoModelForSeq2SeqLM, GenerationConfig
import re
import time

tokenizer = AutoTokenizer.from_pretrained("GIGAParviz/T5_fa_law_chatbot")
model = AutoModelForSeq2SeqLM.from_pretrained("GIGAParviz/T5_fa_law_chatbot").to("cpu")

def clear_text(text):
    if isinstance(text, str):
        text = re.sub("[^آ-ی۰-۹]+", " ", text)
    else:
        text = str(text)
    return text

def generate_response(message, chat_history):
    question = clear_text(message)
    prompt = f"Answer The Question in farsi: {question} Answer: "
    
    tokenized_test_text = tokenizer(prompt, return_tensors='pt').input_ids.to("cpu")
    
    generation_config = GenerationConfig(
        max_new_tokens=128,
        do_sample=True,
        top_k=50,  
        top_p=0.95, 
        temperature=0.7,  
        repetition_penalty=1.2 
    )

    test_output = model.generate(tokenized_test_text, generation_config=generation_config)[0]
    bot_message = tokenizer.decode(test_output, skip_special_tokens=True)
    
    for i in range(0, len(bot_message), 10):
        yield chat_history + [(message, bot_message[:i + 10])]
        time.sleep(0.1)
    
    yield chat_history + [(message, bot_message)]

with gr.Blocks() as demo:
    gr.Markdown("<h1 style='text-align: center;'>💬Persian legal chatbot </h1><p style='text-align: center;'>made by A.M.Parviz/p>")
    
    chatbot = gr.Chatbot(label="جواب")
    msg = gr.Textbox(label="ورودی", placeholder="...سوال حقوقی خودتون رو وارد کنید", lines=1)
    
    msg.submit(generate_response, [msg, chatbot], chatbot)
    
    clear = gr.ClearButton([msg, chatbot])

demo.launch()