import gradio as gr
from transformers import pipeline, AutoTokenizer
import torch
import spaces

#model_name = "DISLab/SummLlama3.2-3B"
model_name = "DISLab/SummLlama3.1-8B"
print(f"Loading model: {model_name}")

tokenizer = AutoTokenizer.from_pretrained(model_name)
pipe = pipeline(
    "text-generation",
    model=model_name,
    tokenizer=tokenizer,
    device_map="auto",
    torch_dtype=torch.float16
)

print("Model loaded successfully!")

def format_chat_template(instruction, document):
    enhanced_instruction = f"{instruction} If the input contains timestamps, preserve them and note when key topics were discussed."
    
    row_json = [{
        "role": "user", 
        "content": f"Below is an instruction that describes a task. Write a response that appropriately completes the request.\n\n### Instruction:\n{enhanced_instruction}\n\n### Input:\n{document}\n\n### Response:\n"
    }]
    return tokenizer.apply_chat_template(row_json, tokenize=False, add_generation_prompt=False)

@spaces.GPU
def summarize(instruction, text):
    try:
        formatted_input = format_chat_template(instruction, text)
        
        output = pipe(
            formatted_input,
            max_new_tokens=3000,
            do_sample=True,
            temperature=0.7,
            top_p=0.9,
            return_full_text=False
        )
        
        summary = output[0]['generated_text'].strip()
        return summary
        
    except Exception as e:
        return f"Error: {str(e)}"

demo = gr.Interface(
    fn=summarize,
    inputs=[
        gr.Textbox(
            lines=3,
            value="Please provide a concise summary of this meeting transcript.",
            label="Custom Instruction",
            placeholder="Enter your instruction here..."
        ),
        gr.Textbox(
            lines=10, 
            placeholder="Enter text or meeting transcript to summarize...",
            label="Document/Transcript"
        )
    ],
    outputs=gr.Textbox(
        label="Summary",
        lines=8
    ),
    title="SummLlama3.2-3B Summarization",
    description="Test the DISLab/SummLlama3.2-3B model with customizable instructions. Timestamps are automatically preserved when present in the input.",
    examples=[
        ["Please provide a concise summary of this meeting transcript."],
        ["Summarize the key technical points discussed."],
        ["Extract action items and decisions, including who is responsible."],
        ["Provide a brief summary highlighting main topics."],
        ["Focus on technical details and implementation discussions."],
        ["Create a bullet-point summary of key takeaways."]
    ],
    allow_flagging="never"
)

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