Update app.py

app.py

@@ -1,71 +1,33 @@
-import torch
-from transformers import AutoModelForCausalLM, AutoTokenizer
 import gradio as gr
+from transformers import AutoTokenizer, AutoModelForCausalLM
+import torch
 
-# Load model dan tokenizer
-model_id = "cahya/gpt2-small-indonesian-522M"
-tokenizer = AutoTokenizer.from_pretrained(model_id)
-model = AutoModelForCausalLM.from_pretrained(model_id, torch_dtype=torch.float16 if torch.cuda.is_available() else torch.float32)
-model.to("cuda" if torch.cuda.is_available() else "cpu")
-model.eval()
-
-# Baca prompt awal
-try:
-    with open("prompt.txt", "r", encoding="utf-8") as f:
-        base_prompt = f.read()
-except FileNotFoundError:
-    base_prompt = ""
-
-# Fungsi buat ngegabungin chat history
-def generate_prompt(message, chat_history):
-    full_prompt = base_prompt
-    for user_msg, ai_msg in chat_history:
-        full_prompt += f"Arya: {user_msg}\nCici: {ai_msg}\n"
-    full_prompt += f"Arya: {message}\nCici:"
-    return full_prompt
+# Load GPT-2 Bahasa Indonesia (contoh: IndoGPT2 dari cahya)
+model_name = "cahya/gpt2-small-indonesian-522M"
+tokenizer = AutoTokenizer.from_pretrained(model_name)
+model = AutoModelForCausalLM.from_pretrained(model_name)
 
-# Fungsi buat ngerespon input user
-def predict(message, chat_history):
-    prompt = generate_prompt(message, chat_history)
-    inputs = tokenizer(prompt, return_tensors="pt", truncation=True, max_length=1024).to(model.device)
-    
+def generate_response(prompt):
+    input_ids = tokenizer.encode(prompt, return_tensors="pt")
     with torch.no_grad():
-        outputs = model.generate(
-            **inputs,
-            max_new_tokens=128,
+        output = model.generate(
+            input_ids,
+            max_length=100,
             do_sample=True,
-            temperature=0.9,
+            top_k=50,
             top_p=0.95,
+            temperature=0.9,
             pad_token_id=tokenizer.eos_token_id,
-            eos_token_id=tokenizer.eos_token_id,
         )
-    
-    output_text = tokenizer.decode(outputs[0], skip_special_tokens=True)
-    
-    # Ambil jawaban terakhir setelah "Cici:"
-    if "Cici:" in output_text:
-        answer = output_text.split("Cici:")[-1].strip()
-    else:
-        answer = "Maaf ya, Cici bingung jawabnya 😢"
-    
-    chat_history.append((message, answer))
-    return answer, chat_history
-
-# Komponen Gradio
-chatbot = gr.Chatbot()
-with gr.Blocks(css=".gradio-container {background-color: #fefefe}") as demo:
-    gr.Markdown("<h1 style='text-align: center;'>🩷 Cici Chatbot Indo 😳🤭</h1>")
-    with gr.Row():
-        with gr.Column():
-            message = gr.Textbox(label="Ketik di sini sayang~ 😚")
-            clear = gr.Button("🧹 Bersihin Chat")
-        with gr.Column():
-            output = chatbot
-        
-    state = gr.State([])
+    return tokenizer.decode(output[0], skip_special_tokens=True)
 
-    message.submit(predict, [message, state], [output, state])
-    clear.click(lambda: ([], []), None, [output, state])
+# Gradio UI
+iface = gr.Interface(
+    fn=generate_response,
+    inputs=gr.Textbox(lines=4, placeholder="Ketik sesuatu yang nakal... 🤭"),
+    outputs="text",
+    title="🧠 AI Nakal Bahasa Indonesia",
+    description="Powered by GPT-2 Indo | Coba ketik rayuan atau pertanyaan nakal 😳💕"
+)
 
-# Launch
-demo.launch()
+iface.launch()