feat: allow usage of fine-tuned model for chart generation

Dockerfile

@@ -7,6 +7,10 @@ COPY . .
 RUN pip install --upgrade pip
 RUN pip install -r requirements.txt
 
+# Ensure transformers and huggingface-hub cache is shared (optional, for faster cold starts)
+ENV TRANSFORMERS_CACHE=/code/.cache/huggingface/transformers
+ENV HF_HOME=/code/.cache/huggingface
+
 EXPOSE 7860
 
 CMD ["python", "app.py"]

README.md

@@ -7,4 +7,5 @@ sdk: gradio
 sdk_version: "4.16.0"
 app_file: app.py
 pinned: false
+# Model: ArchCoder/fine-tuned-bart-large is used for all BART-based inference
 ---

llm_agent.py

@@ -20,7 +20,8 @@ class LLM_Agent:
         self.chart_generator = ChartGenerator(self.data_processor.data)
         self.image_verifier = ImageVerifier()
 
-        model_path = os.path.join(os.path.dirname(__file__), "fine-tuned-bart-large")
+        # Use Hugging Face Hub model path for fine-tuned model
+        model_path = "ArchCoder/fine-tuned-bart-large"
         self.query_tokenizer = AutoTokenizer.from_pretrained(model_path)
         self.query_model = AutoModelForSeq2SeqLM.from_pretrained(model_path)
 

requirements.txt

@@ -64,7 +64,7 @@ tokenizers
 torch
 torchvision
 tqdm
-transformers
+transformers>=4.36.0
 typing_extensions
 tzdata
 urllib3