Update app/model.py

app/model.py

@@ -1,23 +1,25 @@
-from transformers import AutoTokenizer, AutoModelForCausalLM
-import torch
 import os
+import torch
+from transformers import AutoModelForCausalLM, AutoTokenizer
 
-MODEL_NAME = "ayureasehealthcare/llama3-ayurveda-text-v4"
+# Set cache directory early
+CACHE_DIR = "/data/cache"
+os.environ["TRANSFORMERS_CACHE"] = CACHE_DIR
+os.environ["HF_HOME"] = CACHE_DIR
 
-# Set writable HF cache directory
-os.environ["TRANSFORMERS_CACHE"] = "/app/cache/huggingface"
-os.environ["HF_HOME"] = "/app/cache/huggingface"
+MODEL_NAME = "ayureasehealthcare/llama3-ayurveda-text-v4"
 
 def load_model():
-    tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)
+    tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME, cache_dir=CACHE_DIR)
     model = AutoModelForCausalLM.from_pretrained(
         MODEL_NAME,
         torch_dtype=torch.float16 if torch.cuda.is_available() else torch.float32,
-        device_map="auto"
+        device_map="auto",
+        cache_dir=CACHE_DIR
     )
     return model, tokenizer
 
-def generate_response(model, tokenizer, prompt, max_new_tokens=512, temperature=0.7, top_p=0.95):
+def generate_response(model, tokenizer, prompt, max_new_tokens=512, temperature=0.7, top_p=0.9):
     inputs = tokenizer(prompt, return_tensors="pt").to(model.device)
     outputs = model.generate(
         **inputs,
@@ -26,5 +28,4 @@ def generate_response(model, tokenizer, prompt, max_new_tokens=512, temperature=
         top_p=top_p,
         do_sample=True
     )
-    result = tokenizer.decode(outputs[0], skip_special_tokens=True)
-    return result[len(prompt):].strip()
+    return tokenizer.decode(outputs[0], skip_special_tokens=True)