# Use a pipeline as a high-level helper
from transformers import pipeline
import re
import gradio as gr

pipe = pipeline("text-generation", model="microsoft/Phi-4-mini-instruct", trust_remote_code=True)

def extract_reply(generated_text: str) -> str:
    # Remove <think>...</think> block if present
    cleaned = re.sub(r"<think>.*?</think>", "", generated_text, flags=re.DOTALL).strip()

    # If <think> was not used as XML-style tag, remove anything before the final message
    if "</think>" in generated_text and cleaned != generated_text:
        return cleaned

    # If there's no <think>, try to get last paragraph (heuristic)
    lines = [line.strip() for line in cleaned.split("\n") if line.strip()]
    return lines[-1] if lines else cleaned



def simple_bot_with_history(msg, history):
    messages = history +  [
        {"role": "user", "content": msg},
    ]


    result = pipe( messages, )
    generated_text = result[-1]['generated_text'][-1]['content']
    generated_text = extract_reply(generated_text)
    return generated_text

demo = gr.ChatInterface(
    fn=simple_bot_with_history,
    type="messages"
)
demo.launch()