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

# Setup model dan tokenizer
torch.random.manual_seed(0)
model = AutoModelForCausalLM.from_pretrained(
    "microsoft/Phi-3-mini-128k-instruct",
    device_map="cpu",  # Gunakan 'cpu' jika tidak ada GPU
    torch_dtype="auto",
    trust_remote_code=True,
    attn_implementation="eager"  # Menggunakan eager untuk menghindari masalah flash-attention
)

tokenizer = AutoTokenizer.from_pretrained("microsoft/Phi-3-mini-128k-instruct")

# Pipeline untuk text-generation
pipe = pipeline(
    "text-generation",
    model=model,
    tokenizer=tokenizer,
)

# Fungsi untuk menghasilkan respons
def generate_response(input_text):
    messages = [
        {"role": "system", "content": "You are a helpful AI assistant."},
        {"role": "user", "content": input_text}
    ]
    
    generation_args = {
        "max_new_tokens": 500,
        "return_full_text": False,
        "temperature": 0.7,  # Bisa disesuaikan untuk variasi output
        "do_sample": True,   # Mengaktifkan sampling untuk variasi output
    }

    output = pipe(messages, **generation_args)
    return output[0]['generated_text']

# Membuat antarmuka menggunakan Gradio
iface = gr.Interface(
    fn=generate_response,             # Fungsi untuk menangani input
    inputs=gr.Textbox(label="Ask me anything!", placeholder="Tanyakan sesuatu..."),  # Input teks
    outputs=gr.Textbox(label="AI Response"),  # Output teks dari AI
    title="AI Chatbot Assistant",      # Judul aplikasi
    description="Tanya apapun, saya siap membantu!",  # Deskripsi aplikasi
)

# Menjalankan antarmuka
iface.launch()