hugging2021 commited on
Commit
a6fb29f
·
verified ·
1 Parent(s): 92d2564

Update rag_server.py

Browse files
Files changed (1) hide show
  1. rag_server.py +12 -12
rag_server.py CHANGED
@@ -13,22 +13,22 @@ from transformers import AutoModel
13
  import streamlit as st
14
 
15
  # --- Konfiguration ---
16
- os.environ["HF_HOME"] = "/app/cache" # Specify cache path
17
- HF_API_TOKEN = os.environ.get("HF_API_TOKEN") # Read token from environment variable
18
  MODEL_NAME = "dannyk97/mistral-screenplay-model"
 
19
 
20
  # --- Hilfsfunktionen ---
21
 
22
  def query_huggingface_inference_endpoints(prompt):
23
  """
24
- Sends a request to the Hugging Face Inference API.
25
  """
26
  try:
27
  client = InferenceClient(token=HF_API_TOKEN)
28
  result = client.text_generation(prompt, model=MODEL_NAME)
29
  return result
30
  except Exception as e:
31
- return f"Error in query_huggingface_inference_endpoints: {e}"
32
 
33
  # Function to download PDF from Google Drive
34
  def download_pdf_from_drive(drive_link):
@@ -56,16 +56,16 @@ def chunk_text(text, chunk_size=500, chunk_overlap=50):
56
  return text_splitter.split_text(text)
57
 
58
  # Function to create embeddings and store in FAISS
59
- def create_embeddings_and_store(chunks):
60
  try:
61
- embeddings = HuggingFaceEmbeddings(model_name="sentence-transformers/all-MiniLM-L6-v2")
62
  vector_db = FAISS.from_texts(chunks, embedding=embeddings)
63
  return vector_db
64
  except Exception as e:
65
- print(f"Error creating embeddings: {e}")
66
- print("Using dummy embeddings to proceed (functionality will be limited).")
67
- # Fallback to a simpler embedding model (but this might not work well)
68
- vector_db = FAISS.from_texts(["fallback text"], HuggingFaceEmbeddings(model_name="all-mpnet-base-v2")) #Ggf mit "" ersetzen, falls die Implementierung nicht passt.
69
  return vector_db
70
 
71
  # Function to query the vector database and interact with Hugging Face Inference API
@@ -77,7 +77,7 @@ def query_vector_db(query, vector_db):
77
  # Interact with the Text Generation API
78
  prompt = f"Nutze diesen Kontext um die Frage zu beantworten: {context}\nFrage: {query}"
79
  try:
80
- output = query_huggingface_inference_endpoints(prompt) #Keine Modelangabe mehr
81
  return output
82
  except Exception as e:
83
  return f"FEHLER: {str(e)}"
@@ -116,7 +116,7 @@ for link in drive_links:
116
 
117
  if all_chunks:
118
  # Generate embeddings and store in FAISS
119
- vector_db = create_embeddings_and_store(all_chunks)
120
  st.write("Embeddings Generated and Stored Successfully!")
121
 
122
  # User query input
 
13
  import streamlit as st
14
 
15
  # --- Konfiguration ---
16
+ HF_API_TOKEN = os.environ.get("HF_API_TOKEN") # Lesen Sie den Token aus der Umgebungsvariable
 
17
  MODEL_NAME = "dannyk97/mistral-screenplay-model"
18
+ HF_CACHE_DIR = os.environ.get("HF_CACHE_DIR", "/app/cache") #Falls ein Fehler Auftritt, wird der Ordner auf /app/cache gesetzt.
19
 
20
  # --- Hilfsfunktionen ---
21
 
22
  def query_huggingface_inference_endpoints(prompt):
23
  """
24
+ Stellt eine Anfrage an die Hugging Face Inference API.
25
  """
26
  try:
27
  client = InferenceClient(token=HF_API_TOKEN)
28
  result = client.text_generation(prompt, model=MODEL_NAME)
29
  return result
30
  except Exception as e:
31
+ return f"Fehler bei der Anfrage an Hugging Face API: {e}"
32
 
33
  # Function to download PDF from Google Drive
34
  def download_pdf_from_drive(drive_link):
 
56
  return text_splitter.split_text(text)
57
 
58
  # Function to create embeddings and store in FAISS
59
+ def create_embeddings_and_store(chunks, cache_folder=HF_CACHE_DIR):
60
  try:
61
+ embeddings = HuggingFaceEmbeddings(model_name="sentence-transformers/all-MiniLM-L6-v2", cache_folder=cache_folder)
62
  vector_db = FAISS.from_texts(chunks, embedding=embeddings)
63
  return vector_db
64
  except Exception as e:
65
+ print(f" Fehler beim Erstellen der Embeddings: {e}")
66
+ print("Verwende Dummy Embeddings, um fortzufahren (Funktionen sind eingeschränkt).")
67
+ # Verwenden Sie eine einfachere Fallback Lösung
68
+ vector_db = FAISS.from_texts(["fallback text"], HuggingFaceEmbeddings(model_name="all-mpnet-base-v2", cache_folder=cache_folder))
69
  return vector_db
70
 
71
  # Function to query the vector database and interact with Hugging Face Inference API
 
77
  # Interact with the Text Generation API
78
  prompt = f"Nutze diesen Kontext um die Frage zu beantworten: {context}\nFrage: {query}"
79
  try:
80
+ output = query_huggingface_inference_endpoints(prompt)
81
  return output
82
  except Exception as e:
83
  return f"FEHLER: {str(e)}"
 
116
 
117
  if all_chunks:
118
  # Generate embeddings and store in FAISS
119
+ vector_db = create_embeddings_and_store(all_chunks, cache_folder=HF_CACHE_DIR)
120
  st.write("Embeddings Generated and Stored Successfully!")
121
 
122
  # User query input