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

model_id = "DLMveloper/Solade-DLM-7B-4"
token = os.getenv("HF_TOKEN")

print("Loading Model")
tokenizer = AutoTokenizer.from_pretrained(model_id, token=token)
model = AutoModelForCausalLM.from_pretrained(
    model_id,
    torch_dtype=torch.float16,
    device_map="auto",
    token=token
)
print("Model Complete")

@spaces.GPU
def generate_api(prompt: str, max_new_tokens: int, temperature: float):
    inputs = tokenizer(prompt, return_tensors="pt").to(model.device)
    
    outputs = model.generate(
        **inputs,
        max_new_tokens=int(max_new_tokens),
        temperature=float(temperature),
        do_sample=True,
        eos_token_id=model.config.eos_token_id
    )
    
    response = tokenizer.decode(outputs[0][inputs.input_ids.shape[1]:], skip_special_tokens=True)
    return response

demo = gr.Interface(
    fn=generate_api,
    inputs=[
        gr.Textbox(label="Prompt"),
        gr.Slider(minimum=16, maximum=2048, value=512, step=1, label="Max New Tokens"),
        gr.Slider(minimum=0.1, maximum=2.0, value=0.7, step=0.05, label="Temperature")
    ],
    outputs=gr.Textbox(label="Response")
)

if __name__ == "__main__":
    demo.launch(server_name="0.0.0.0", server_port=7860)