import gradio as gr
from transformers import AutoModelForCausalLM, AutoTokenizer
import torch
from peft import (
    LoraConfig,
    PeftModel,
    prepare_model_for_kbit_training,
    get_peft_model,
)
model_name = "meta-llama/Llama-3.2-1B"
lora_model_name="Anlam-Lab/Llama-3.2-1B-it-anlamlab-SA-Chatgpt4mini"

tokenizer = AutoTokenizer.from_pretrained(model_name)
model = AutoModelForCausalLM.from_pretrained(model_name, device_map="auto", torch_dtype=torch.float16)

model = PeftModel.from_pretrained(model, lora_model_name)
def generate_response(input_text):
    inputs = tokenizer(input_text, return_tensors="pt").to(model.device)
    
    generation_config = {
        "max_length": 512,
        "temperature": 0.01,
        "do_sample": True,
        "pad_token_id": tokenizer.pad_token_id,
        "eos_token_id": tokenizer.eos_token_id,
    }
    
    with torch.no_grad():
        outputs = model.generate(
            **inputs,
            **generation_config
        )
    
    response = tokenizer.decode(outputs[0])
    return response.split("<|end_header_id|>")[-1].split("<|eot_id|>")[0]

iface = gr.Interface(
    fn=generate_response,
    inputs=gr.Textbox(lines=5, placeholder="Metninizi buraya girin..."),
    outputs=gr.Textbox(lines=5, label="Model Çıktısı"),
    title="Anlam-Lab",
    examples=[
        ["Akıllı saati uzun süre kullandım ve şık tasarımı, harika sağlık takibi özellikleri ve uzun pil ömrüyle çok memnun kaldım."],
        ["Ürünü aldım ama pil ömrü kısa, ekran parlaklığı yetersiz ve sağlık takibi doğru sonuçlar vermedi."],
    ]
)

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