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",
        device_map="auto",
        torch_dtype="auto"
    )
    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(model.device)
    
    outputs = model.generate(
        **inputs,
        max_new_tokens=128,
        temperature=0.7,
        pad_token_id=tokenizer.eos_token_id
    )
    
    full_text = tokenizer.decode(outputs[0], skip_special_tokens=True)
    return full_text.split("### Summary:")[-1].strip()