tsuin-complete / app.py
debisoft's picture
pre-load
7842d10
raw
history blame
911 Bytes
import gradio as gr
import huggingface_hub
import os
import spaces
import torch
from transformers import AutoTokenizer, AutoModelForCausalLM
huggingface_hub.login(os.getenv('HF_TOKEN'))
tokenizer = AutoTokenizer.from_pretrained("deepseek-ai/DeepSeek-R1-Distill-Qwen-7B")
model = AutoModelForCausalLM.from_pretrained("deepseek-ai/DeepSeek-R1-Distill-Qwen-7B").to(device)
cuda_device = torch.device("cuda")
cpu_device = torch.device("cpu")
@spaces.GPU
def sentience_check():
model.to(cuda_device)
inputs = tokenizer("Are you sentient?", return_tensors="pt").to(device)
with torch.no_grad():
outputs = model.generate(
**inputs, max_new_tokens=128, pad_token_id = tokenizer.eos_token_id
)
model.to(cpu_device)
return tokenizer.decode(outputs[0], skip_special_tokens=True)
demo = gr.Interface(fn=sentience_check, inputs=None, outputs=gr.Text())
demo.launch()