Spaces:
Running
Running
File size: 581 Bytes
4a69777 890ddfa 582ae31 890ddfa 4a69777 39decae 4a69777 39decae 3b2f6fb 39decae 3b2f6fb 39decae 4a69777 582ae31 4a69777 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 |
import streamlit as st
from transformers import AutoModelForCausalLM, AutoTokenizer
import torch
# Load model and tokenizer
@st.cache_resource()
def load_model():
model_name = "deepseek-ai/deepseek-coder-1.3b-instruct" # Using smaller model
tokenizer = AutoTokenizer.from_pretrained(model_name)
model = AutoModelForCausalLM.from_pretrained(
model_name,
torch_dtype=torch.float32, # Use float32 for CPU (float16 needs GPU)
device_map={"": "cpu"} # Force model to use CPU
)
return model, tokenizer
model, tokenizer = load_model()
|