inference.py

from transformers import AutoModelForCausalLM, AutoTokenizer, pipeline
import torch
import os
import spaces 

# Load the model and tokenizer
print("Starting to load the model...")
hf_token = os.environ.get("HF_TOKEN")
model_id = "google/gemma-2-2b-it"

model = AutoModelForCausalLM.from_pretrained(model_id, device_map="auto", torch_dtype=torch.float16, token=hf_token)
tokenizer = AutoTokenizer.from_pretrained(model_id, token=hf_token)
pipe = pipeline("text-generation", model=model, tokenizer=tokenizer)
print("Model loaded successfully!")

# Settings for the model
TEMPERATURE = 0.7
TOP_P = 0.9
MAX_NEW_TOKENS = 60

@spaces.GPU
def generate_excuse(prompt):
    full_prompt = f"""You are an Advanced Excuse Generator AI, designed to create concise, unexpected, creative, and logically sound excuses for various situations. Your task is to generate an excuse that is:

0. Concise
1. Highly unexpected and original
2. Logically valid and difficult to disprove
3. Detailed enough to be believable
4. Tailored to the specific situation provided
5. Slightly absurd, but not so outlandish as to be immediately dismissed

Consider various angles such as:
- Technological glitches or malfunctions
- Obscure natural phenomena
- Little-known historical events or cultural practices
- Rare medical conditions or psychological phenomena
- Unlikely coincidences or chain reactions

For the following situation, generate a concise, creative, and logically sound excuse and only output the excuse do not add anything else:

{prompt}

Excuse:
"""
    inputs = tokenizer(full_prompt, return_tensors="pt").to(model.device)

    with torch.no_grad():
        outputs = model.generate(
            **inputs, 
            max_new_tokens=MAX_NEW_TOKENS, 
            temperature=TEMPERATURE, 
            top_p=TOP_P,
            pad_token_id=tokenizer.eos_token_id, 
            do_sample=True
        )

    excuse = tokenizer.decode(outputs[0], skip_special_tokens=True)
    return excuse.split("Excuse:")[-1].strip(), {"model_id": model_id, "TEMPERATURE": TEMPERATURE, "TOP_P": TOP_P}