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

# Set an environment variable
HF_TOKEN = os.environ.get("HF_TOKEN", None)

# Load the tokenizer and model
tokenizer = AutoTokenizer.from_pretrained("meta-llama/Llama-3.2-1B-Instruct")
model = AutoModelForCausalLM.from_pretrained("meta-llama/Llama-3.2-1B-Instruct", device_map="auto")

def generate_summary(text: str, temperature: float, max_new_tokens: int) -> str:
    """
    Generate a single-line summary from the input text using the llama3-8b model.
    Args:
        text (str): The input text to summarize.
        temperature (float): The temperature for generating the response.
        max_new_tokens (int): The maximum number of new tokens to generate.
    Returns:
        str: The generated summary in a single line.
    """
    input_ids = tokenizer.encode(text, return_tensors="pt").to(model.device)

    output_ids = model.generate(
        input_ids=input_ids,
        max_new_tokens=max_new_tokens,
        do_sample=True,
        temperature=temperature,
    )

    summary = tokenizer.decode(output_ids[0], skip_special_tokens=True)
    # Convert to a single line
    return " ".join(summary.split())

def summarize_file(file_path, temperature: float, max_new_tokens: int) -> str:
    """
    Summarize the content of an uploaded file into a single line.
    Args:
        file_path (str): The path of the uploaded file.
        temperature (float): The temperature for generating the response.
        max_new_tokens (int): The maximum number of new tokens to generate.
    Returns:
        str: The generated summary of the file's content in a single line.
    """
    # Read file content
    with open(file_path, 'r') as f:
        text = f.read()
        
    # Generate summary
    return generate_summary(text, temperature, max_new_tokens)

# Gradio block for text summarization
with gr.Blocks() as demo:
    gr.Markdown("<h1>Text Summarization Application</h1>")
    with gr.Row():
        with gr.Column():
            text_input = gr.Textbox(lines=10, label="Input Text", placeholder="Enter text here...")
            file_input = gr.File(label="Upload Text File", file_count="single", type="filepath")
            temperature = gr.Slider(minimum=0, maximum=1, step=0.1, value=0.7, label="Temperature")
            max_tokens = gr.Slider(minimum=10, maximum=512, step=1, value=150, label="Max New Tokens")
            submit_button = gr.Button("Generate Summary")
        with gr.Column():
            summary_output = gr.Textbox(lines=1, label="Summary", interactive=False)

    # Link button to generate summary from text input
    submit_button.click(
        fn=generate_summary,
        inputs=[text_input, temperature, max_tokens],
        outputs=summary_output
    )

    # Link file upload to summary generation
    file_input.change(
        fn=summarize_file,
        inputs=[file_input, temperature, max_tokens],
        outputs=summary_output
    )

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