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

# Load the model and tokenizer
model_name = "Equall/Saul-7B-Instruct-v1"  # This is the model you're using
model = AutoModelForCausalLM.from_pretrained(model_name)
tokenizer = AutoTokenizer.from_pretrained(model_name)

# Define the prediction function
def generate_response(prompt):
    inputs = tokenizer(f"<s>[INST] Vi ste pravni stručnjak specijaliziran za hrvatsko pravo... {prompt} [/INST]</s>", return_tensors="pt")
    with torch.no_grad():
        output = model.generate(inputs["input_ids"], max_length=500, temperature=0.7, top_p=0.9)
    response = tokenizer.decode(output[0], skip_special_tokens=True)
    return response.strip()

# Create the Gradio interface
iface = gr.Interface(fn=generate_response, inputs="text", outputs="text", title="Croatia Legal Assistant", description="Ask legal questions related to Croatian law.")

# Launch the Gradio interface
iface.launch()