acecalisto3 commited on
Commit
b7cffa1
·
1 Parent(s): ac4d529
Files changed (1) hide show
  1. app.py +24 -2
app.py CHANGED
@@ -46,8 +46,30 @@ model_loaders = {
46
 
47
  model_option = st.selectbox("Select a Model", list(model_loaders.keys()))
48
 
49
- # Load the selected model
50
- model = model_loaders[model_option]()
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
51
 
52
  HUGGINGFACE_TOKEN = os.getenv("HUGGINGFACE_TOKEN")
53
 
 
46
 
47
  model_option = st.selectbox("Select a Model", list(model_loaders.keys()))
48
 
49
+ def load_model(model_name: str):
50
+ """
51
+ Loads the specified model and tokenizer.
52
+ """
53
+ try:
54
+ tokenizer = AutoTokenizer.from_pretrained(model_name, use_fast=False, legacy=False)
55
+ model = AutoModelForCausalLM.from_pretrained(model_name)
56
+ # This should be inside the try block
57
+ max_supported_length = 2048 # Get this from the model config
58
+ openllama_pipeline = pipeline(
59
+ "text-generation",
60
+ model=model,
61
+ tokenizer=tokenizer,
62
+ truncation=True,
63
+ max_length=max_supported_length,
64
+ temperature=0.7,
65
+ top_p=0.95,
66
+ device=0 if torch.cuda.is_available() else -1,
67
+ )
68
+ logging.info(f"{model_name} loaded successfully.")
69
+ return openllama_pipeline
70
+ except Exception as e:
71
+ logging.error(f"Error loading {model_name} model: {e}")
72
+ return None
73
 
74
  HUGGINGFACE_TOKEN = os.getenv("HUGGINGFACE_TOKEN")
75