#llama.py
from transformers import AutoTokenizer, AutoModelForCausalLM

PROMPT_TEMPLATE = """
You are an expert summarizer. Generate concise German summaries based on the email content,
following this structure: 'Der Kunde ... und erwartet ...'. The summaries need to be short. Here are some examples:
- Der Kunde übermittelt den Kontoauszug mit einer Abbuchung von 103,22€ und erwartet die Bestätigung der Abbuchung.
- Der Kunde möchte die Zahlungsart von Überweisung auf Lastschrift ändern und erwartet die Änderung der Zahlungsart durch die Unternehmensvertretung.
- Der Kunde übermittelt fehlende Angaben wie Übergabedatum und Zählerstand und erwartet die Verarbeitung der bereitgestellten Informationen.
- Der Kunde teilt die Beendigung des Gasbezugs aufgrund der Installation einer Wärmepumpe mit und erwartet den Abschluss des Gasvertrags und Bestätigung.
- Der Kunde sendet Daten und Papiere zur Entsperrung des Stromzählers.
- Der Kunde bittet um Korrektur der berechneten Abschlagszahlung für August 2023 und erwartet die Überprüfung und Anpassung der Zahlung.
- Der Kunde fragt nach der Abrechnung für einen beendeten Vertrag und erwartet die Erstellung und Zusendung der Abrechnung.
- Der Kunde bittet um Stellungnahme oder Korrektur, da der E.ON Plus Rabatt nicht berücksichtigt wurde und erwartet die Überprüfung und Korrektur der Rechnungen.
- Der Kunde sendet ein unterschriebenes Formular zurück und bittet um Kenntnisnahme und erwartet die Verarbeitung des Formulars.
- Der Kunde bittet um eine Ratenzahlung von 30 Euro monatlich für die letzte Stromrechnung und erwartet die Vereinbarung der Ratenzahlung.

### Email:
{}

### Summary:
"""
def load_llama_model():
    """Load Llama model and tokenizer with optimized settings"""
    tokenizer = AutoTokenizer.from_pretrained("Walid777/llama3-8b-emails-summarization")
    model = AutoModelForCausalLM.from_pretrained( "Walid777/llama3-8b-emails-summarization")
    return model, tokenizer

def generate_llama_summary(email, model, tokenizer, prompt_template):
    """Generate summary using structured prompt template"""
    formatted_prompt = prompt_template.format(email, "")
    
    inputs = tokenizer(
        [formatted_prompt],
        return_tensors="pt"
    ).to("cuda")
    
    outputs = model.generate(
        **inputs,
        max_new_tokens=128
    )
    
    summary = tokenizer.decode(outputs[0], skip_special_tokens=True)
    if "### Summary:" in summary:
        summary = summary.split("### Summary:")[-1].strip()
    else:
        summary = "Error: Could not extract summary"
    if "### Email:" in summary:
        summary = summary.split("### Email:")[0].strip()
    return summary