Spaces:
Runtime error
Runtime error
File size: 781 Bytes
c534dde 4b78593 c534dde 5f9ff8f c534dde 4b78593 c534dde |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 |
import gradio as gr
from transformers import AutoModelForCausalLM, AutoTokenizer
import safetensors.torch
import spaces
# Load base model and tokenizer
base_model_name = "AIDC-AI/Marco-o1"
tokenizer = AutoTokenizer.from_pretrained(base_model_name)
model = AutoModelForCausalLM.from_pretrained(base_model_name)
# Load your lightweight safetensor
safetensor_path = "MegaTronX/Odyssey-SelectolaxQLoRA"
model.load_state_dict(safetensors.torch.load_file(safetensor_path), strict=False)
@spaces.GPU
def predict(text):
inputs = tokenizer.encode(text, return_tensors="pt")
outputs = model.generate(inputs)
prediction = tokenizer.decode(outputs[0], skip_special_tokens=True)
return prediction
demo = gr.Interface(fn=predict, inputs="text", outputs="text")
demo.launch()
|