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

model_id = "forti2026/gemma-3-1b-chatbot-skripsi"

print(f"Sedang mendownload model baru: {model_id}")

tokenizer = AutoTokenizer.from_pretrained(model_id)

model = AutoModelForCausalLM.from_pretrained(
    model_id,
    torch_dtype=torch.float32, 
    low_cpu_mem_usage=True
)

def chat_logic(message):
    input_text = f"<start_of_turn>user\n{message}<end_of_turn>\n<start_of_turn>model\n"
    
    inputs = tokenizer(input_text, return_tensors="pt")
    
    # Generate
    outputs = model.generate(
        **inputs, 
        max_new_tokens=250,
        do_sample=True,
        temperature=0.7,
        top_k=50,
        top_p=0.95
    )
    
    response = tokenizer.decode(outputs[0], skip_special_tokens=True)
    clean_response = response.split("model\n")[-1].strip()
    return clean_response

iface = gr.Interface(fn=chat_logic, inputs="text", outputs="text")
iface.launch()