from unsloth import FastLanguageModel
import gradio as gr
import torch

# Charger le modèle base + adapter (LoRA)
model, tokenizer = FastLanguageModel.from_pretrained(
    model_name = "unsloth/Qwen3-14B",  # Base model utilisé pendant l'entraînement
    adapter_path = "ilyass31/DH-AI-negotiation-assistant",  # Ton repo LoRA
    max_seq_length = 2048,
    load_in_4bit = True,
)

# Fonction de génération de texte
def generate_response(prompt, temperature=0.7, max_new_tokens=512):
    inputs = tokenizer(prompt, return_tensors="pt").to("cuda" if torch.cuda.is_available() else "cpu")
    outputs = model.generate(
        **inputs,
        max_new_tokens=max_new_tokens,
        do_sample=True,
        temperature=temperature,
        pad_token_id=tokenizer.eos_token_id
    )
    return tokenizer.decode(outputs[0], skip_special_tokens=True)

# Interface 
demo = gr.Interface(
    fn=generate_response,
    inputs=[
        gr.Textbox(label="Prompt", placeholder="Enter your negotiation scenario..."),
        gr.Slider(0.1, 1.5, value=0.7, step=0.1, label="Temperature"),
        gr.Slider(32, 2048, value=512, step=32, label="Max New Tokens")
    ],
    outputs=gr.Textbox(label="Response"),
    title="🧠 AI Negotiation Assistant",
    description="Enter a scenario and get a negotiation strategy or analysis using a fine-tuned Qwen3-14B model."
)

if __name__ == "__main__":
    demo.launch()