Update app.py

app.py

@@ -3,7 +3,7 @@ import numpy as np
 import gradio as gr
 import chromadb
 
-from transformers import AutoModel, AutoTokenizer, pipeline
+from transformers import AutoModel, AutoTokenizer, AutoModelForCausalLMpipeline
 import torch
 import chromadb
 
@@ -24,7 +24,7 @@ def get_embedding(text):
     return outputs.last_hidden_state[:, 0, :].numpy().tolist()  # Take CLS token embedding
 
 # Load LLaMA Model (Meta LLaMA 2)
-llama_pipe = pipeline("text-generation", model="meta-llama/Llama-2-7b-chat-hf")
+llama_pipe = pipeline("text-generation", model=AutoModelForCausalLM.from_pretrained("meta-llama/Llama-2-7b-chat-hf"))
 
 # Load a small subset (10,000 rows)
 #dataset = load_dataset("wiki40b", "en", split="train[:1000]")