added Llama3

app.py

@@ -3,8 +3,11 @@ from transformers import pipeline
 from fastai.vision.all import *
 from PIL import Image
 import os
+import io
+import json
 
 access_token = os.getenv("access_token")
+
 # NOTE - we configure docs_url to serve the interactive Docs at the root path
 # of the app. This way, we can use the docs as a landing page for the app on Spaces.
 app = FastAPI(docs_url="/")
@@ -13,6 +16,16 @@ pipe = pipeline("text2text-generation", model="google/flan-t5-small")
 categories = ('Heart', 'Oblong', 'Oval', 'Round', 'Square')
 learn = load_learner('model.pkl')
 
+# Initialize the Code Llama Instruct pipeline (example with 7B model)
+llama_model_id = "meta-llama/CodeLlama-7b-Instruct-hf"
+llama_pipeline = pipeline(
+    "text-generation",
+    model=llama_model_id,
+    model_kwargs={"torch_dtype": torch.bfloat16},
+    device_map="auto",
+    use_auth_token=access_token  # Use the access token for authentication
+)
+
 @app.get("/generate")
 def generate(text: str):
     """
@@ -53,3 +66,36 @@ async def face_analyse(file: UploadFile = File(...)):
     # Assuming categories is a list of category labels
     return dict(zip(categories, map(float, probs)))
 
+@app.post("/extract-frame-details")
+def extract_frame_details(text: str):
+    """
+    Using the Code Llama Instruct pipeline from `transformers`, extract frame
+    details from the given input text. The model used is `meta-llama/CodeLlama-7b-Instruct-hf`.
+    """
+    messages = [
+        {"role": "system", "content": "Please provide details about frames in JSON format."},
+        {"role": "user", "content": text},
+    ]
+
+    terminators = [
+        llama_pipeline.tokenizer.eos_token_id,
+        llama_pipeline.tokenizer.convert_tokens_to_ids("")
+    ]
+
+    outputs = llama_pipeline(
+        messages,
+        max_new_tokens=256,
+        eos_token_id=terminators,
+        do_sample=True,
+        temperature=0.6,
+        top_p=0.9,
+    )
+    
+    generated_text = outputs[0]["generated_text"]
+    
+    try:
+        extracted_info = json.loads(generated_text)
+    except json.JSONDecodeError:
+        return {"error": "Failed to parse the generated text as JSON."}
+
+    return extracted_info