CodeCorrect-AI / app.py
Shushmita's picture
Update app.py
39decae verified
raw
history blame contribute delete
581 Bytes
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()