Pls be fixed

Dockerfile

@@ -8,7 +8,9 @@ RUN pip install --upgrade pip
 RUN pip install -r requirements.txt
 
 RUN mkdir /api/cache
+RUN mkdir /api/models
 RUN chmod a+rwx /api/cache
+RUN chmod a+rwx /api/models
 
 EXPOSE 7860
 

main.py

@@ -1,5 +1,5 @@
 from flask import Flask, request, render_template, jsonify
-from models import blenderbot
+from models import llama3
 
 app = Flask("AI API")
 
@@ -16,7 +16,7 @@ def receive_data():
   data = request.get_json()
   print("Prompt:", data["prompt"])
 
-  generated_text = blenderbot.generate(data["prompt"])
+  generated_text = llama3.generate(data["prompt"])
   
   print("Response:", generated_text)
   

models/llama3.py

@@ -0,0 +1,12 @@
+from transformers import AutoTokenizer, AutoModelForCausalLM
+
+model_name = "meta-llama/Meta-Llama-3-8B"
+
+tokenizer = AutoTokenizer.from_pretrained(model_name)
+model = AutoModelForCausalLM.from_pretrained(model_name)
+
+def generate(input_text):
+  input_ids = tokenizer.encode(input_text, return_tensors="pt")
+  output_ids = model.generate(input_ids, num_beams=5, no_repeat_ngram_size=2)
+  
+  return tokenizer.decode(output_ids[0], skip_special_tokens=True)