import os
import gradio as gr
from transformers import AutoModelForCausalLM, AutoTokenizer

# Mendapatkan token Hugging Face
hf_token = os.getenv("HUGGINGFACE_TOKEN")
if not hf_token:
    raise ValueError("Token Hugging Face tidak ditemukan. Harap atur variabel lingkungan 'HUGGINGFACE_TOKEN'.")

# Memuat model dan tokenizer
model_name = "Erlanggaa/TextChatbot"  # Ganti dengan nama repository model kamu di Hugging Face
try:
    tokenizer = AutoTokenizer.from_pretrained(model_name, use_auth_token=hf_token)
    model = AutoModelForCausalLM.from_pretrained(model_name, use_auth_token=hf_token)
except Exception as e:
    raise ValueError(f"Gagal memuat model atau tokenizer: {str(e)}")

# Fungsi untuk menghasilkan teks
def generate_text(prompt):
    if not prompt or len(prompt.strip()) == 0:
        return "Masukkan prompt yang valid."
    
    inputs = tokenizer(prompt, return_tensors="pt")
    
    # Tambahkan parameter yang mengontrol hasil respons
    outputs = model.generate(
        inputs['input_ids'], 
        max_length=500,  # Panjang maksimum respons
        min_length=30,   # Panjang minimum untuk memberikan jawaban lebih baik
        temperature=0.7,  # Mengontrol kreativitas model
        top_k=50,         # Membatasi pilihan model pada 50 token teratas
        top_p=0.9,        # Probabilitas kumulatif untuk sampling
        repetition_penalty=1.2  # Mengurangi pengulangan kata
    )
    
    # Decoding hasil keluaran
    result = tokenizer.decode(outputs[0], skip_special_tokens=True)
    
    # Optional: Post-processing untuk membersihkan hasil
    return result.strip()

# Membuat antarmuka Gradio
interface = gr.Interface(
    fn=generate_text,
    inputs="text",
    outputs="text",
    title="LLaMA 3.2-1B Text Generator",
    description="Masukkan prompt, dan model LLaMA 3.2-1B akan menghasilkan teks."
)

if __name__ == "__main__":
    interface.launch()