import gradio as gr
import torch
import spaces
from transformers import AutoModelForCausalLM, AutoTokenizer, TextIteratorStreamer
from threading import Thread

model_id = "DLMveloper/Test"
tokenizer = AutoTokenizer.from_pretrained(model_id)
model = AutoModelForCausalLM.from_pretrained(model_id, torch_dtype=torch.float16).to("cuda")

@spaces.GPU
def predict(message, history):
    # ПРЕМУБЛЕ (Few-Shot) - заставляем модель видеть диалог
    prompt = """This is a conversation between a User and an AI Assistant.
User: Привет!
AI: Привет! Чем я могу тебе помочь сегодня?
User: Как дела?
AI: Я всего лишь искусственный интеллект, но я готов помочь тебе с вопросами.
"""
    # Добавляем историю
    for human, bot in history:
        prompt += f"User: {human}\nAI: {bot}\n"
    
    # Добавляем текущее сообщение
    prompt += f"User: {message}\nAI:"
    
    inputs = tokenizer(prompt, return_tensors="pt").to("cuda")
    streamer = TextIteratorStreamer(tokenizer, skip_prompt=True, skip_special_tokens=True)
    
    # Очень строгие настройки генерации
    generation_kwargs = dict(
        input_ids=inputs["input_ids"],
        attention_mask=inputs["attention_mask"],
        streamer=streamer,
        max_new_tokens=100,
        do_sample=True,
        temperature=0.4,       # Ставим меньше, чтобы не "фантазировала"
        repetition_penalty=1.5, # Штраф за повторы (1.5 - очень строго)
        top_k=50
    )
    
    thread = Thread(target=model.generate, kwargs=generation_kwargs)
    thread.start()
    
    partial_text = ""
    for new_text in streamer:
        # Если модель начала генерировать бред, мы можем просто остановить её
        if "!!!!!!!!" in new_text:
            break
        partial_text += new_text
        yield partial_text
        
gr.ChatInterface(predict).launch()