cache

.gitignore

@@ -1 +1 @@
-hf_cache
+.hf_cache

app.py

@@ -5,16 +5,18 @@ from transformers import AutoTokenizer, AutoConfig, AutoModelForCausalLM, pipeli
 import gradio as gr
 
 # Load Goldfish model for Dhivehi
-language = 'Dhivehi'
+model_name = 'div_thaa_full'
+HF_CACHE = '.hf_cache'
 
 # Load model
-goldfish_model = 'goldfish-models/div_thaa_full'
-config = AutoConfig.from_pretrained(goldfish_model)
-tokenizer = AutoTokenizer.from_pretrained(goldfish_model)
-model = AutoModelForCausalLM.from_pretrained(goldfish_model, config=config)
+goldfish_model = 'goldfish-models/' + model_name
+config = AutoConfig.from_pretrained(goldfish_model, cache_dir=HF_CACHE)
+tokenizer = AutoTokenizer.from_pretrained(goldfish_model, cache_dir=HF_CACHE)
+model = AutoModelForCausalLM.from_pretrained(goldfish_model, config=config, cache_dir=HF_CACHE)
 if torch.cuda.is_available():
     model = model.cuda()  # Load onto GPU
 
+
 # Create text generation pipeline
 text_generator = pipeline('text-generation', model=model, tokenizer=tokenizer)