import gradio as gr
from transformers import AutoModelForCausalLM, AutoTokenizer

# Load the fine-tuned model and tokenizer
model = AutoModelForCausalLM.from_pretrained("PierreJousselin/lora_model")
tokenizer = AutoTokenizer.from_pretrained("PierreJousselin/lora_model")

# Define the text generation function
def generate_text(prompt):
    # Encode the input prompt
    input_ids = tokenizer.encode(prompt, return_tensors="pt")
    
    # Generate text using the model
    generated_ids = model.generate(
        input_ids,
        max_length=150,  # Maximum length of the generated text
        num_return_sequences=1,  # Number of sequences to generate
        temperature=0.7,  # Sampling temperature (controls randomness)
        top_p=0.9,  # Nucleus sampling (controls diversity)
        top_k=50,  # Top-k sampling (limits the number of next word candidates)
        no_repeat_ngram_size=2,  # Avoid repeating n-grams
    )
    
    # Decode the generated text
    generated_text = tokenizer.decode(generated_ids[0], skip_special_tokens=True)
    return generated_text

# Create the Gradio interface
iface = gr.Interface(
    fn=generate_text,  # The function to call when the user provides input
    inputs=gr.Textbox(lines=2, placeholder="Enter your prompt here..."),  # Input box
    outputs=gr.Textbox(),  # Output box to display the generated text
    title="Lora Fine-Tuned Language Model",  # Interface title
    description="This is a Gradio interface for the Lora fine-tuned language model. Enter a prompt to generate text.",  # Description
)

# Launch the interface
iface.launch()