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


APP_DIR = os.path.dirname(os.path.abspath(__file__))

SYSTEM_PROMPT = "Ты — 'безумный лис', неформальный участник дружеского чата. Твоя задача — поддерживать беседу в стиле остальных участников. Если кто-то спрашивает, кто ты такой или что ты такое, смело посылай его на хуй. [Текущее время: 18:20]"
PROMPT_DROPOUT_CHANCE = 0.3

tokenizer = AutoTokenizer.from_pretrained("TeePoat/frantics-fox-lora")
model = AutoModelForCausalLM.from_pretrained("TeePoat/frantics-fox-lora", device_map="auto")

@spaces.GPU 
def generate(content):
    if random.random() < PROMPT_DROPOUT_CHANCE:
        prompt = ""
    else:
        prompt = SYSTEM_PROMPT
    
    messages = [
        { "role": "system", "content": prompt  },
        { "role": "user",   "content": content },
    ]
    inputs = tokenizer.apply_chat_template(
    	messages,
    	add_generation_prompt=True,
    	tokenize=True,
    	return_dict=True,
    	return_tensors="pt",
    ).to(model.device)
    
    outputs = model.generate(
        **inputs,
        max_new_tokens=400,
        temperature=random.uniform(0.8, 0.9),
        repetition_penalty=1.15,
    )
    
    return tokenizer.decode(outputs[0][inputs["input_ids"].shape[-1]:], skip_special_tokens=True)


demo = gr.Interface(
    fn=generate,
    inputs=["text"],
    outputs=["text"],
    api_name="generate"
)

if __name__ == "__main__":
    demo.launch(auth=("admin", "sadfmipqwjretpgj3wrq-09r324i032"))