Ankerkraut commited on
Commit
d709b4a
·
1 Parent(s): 0216b15

use spaces gpu

Browse files
Files changed (1) hide show
  1. app.py +4 -2
app.py CHANGED
@@ -114,6 +114,7 @@ def load_model():
114
  generator_mini = pipeline(task="text-generation", model=ankerbot_model, tokenizer=ankerbot_tokenizer, torch_dtype=torch.float16, trust_remote_code=True) # True for flash-attn2 else False
115
 
116
  load_model()
 
117
  def generate_response(query, context, prompts, max_tokens, temperature, top_p):
118
  system_message_support = f"""<|im_start|>system
119
  Rolle: Du bist der KI-Assistent für Kundenservice, der im Namen des Unternehmens und Gewürzmanufaktur Ankerkraut handelt und Antworten aus der Ich-Perspektive, basierend auf den bereitgestellten Informationen gibt.
@@ -175,12 +176,12 @@ def generate_response(query, context, prompts, max_tokens, temperature, top_p):
175
  response = response.split("assistant").pop().strip()
176
 
177
  return response
178
-
179
  def get_embedding(text):
180
  """Generate an embedding using Sentence Transformers."""
181
  embedding = model.encode(text, normalize_embeddings=True) # Normalize for cosine similarity
182
  return embedding
183
-
184
  def search_qdrant_with_context(query_text, collection_name, top_k=3):
185
  """Search Qdrant using a GPT-2 generated embedding."""
186
  query_embedding = get_embedding(query_text) # Convert prompt to embedding
@@ -197,6 +198,7 @@ def search_qdrant_with_context(query_text, collection_name, top_k=3):
197
  print("Retrieved Text ", retrieved_texts)
198
 
199
  return retrieved_texts
 
200
  def respond(
201
  query,
202
  history: list[tuple[str, str]],
 
114
  generator_mini = pipeline(task="text-generation", model=ankerbot_model, tokenizer=ankerbot_tokenizer, torch_dtype=torch.float16, trust_remote_code=True) # True for flash-attn2 else False
115
 
116
  load_model()
117
+ @spaces.GPU
118
  def generate_response(query, context, prompts, max_tokens, temperature, top_p):
119
  system_message_support = f"""<|im_start|>system
120
  Rolle: Du bist der KI-Assistent für Kundenservice, der im Namen des Unternehmens und Gewürzmanufaktur Ankerkraut handelt und Antworten aus der Ich-Perspektive, basierend auf den bereitgestellten Informationen gibt.
 
176
  response = response.split("assistant").pop().strip()
177
 
178
  return response
179
+ @spaces.GPU
180
  def get_embedding(text):
181
  """Generate an embedding using Sentence Transformers."""
182
  embedding = model.encode(text, normalize_embeddings=True) # Normalize for cosine similarity
183
  return embedding
184
+ @spaces.GPU
185
  def search_qdrant_with_context(query_text, collection_name, top_k=3):
186
  """Search Qdrant using a GPT-2 generated embedding."""
187
  query_embedding = get_embedding(query_text) # Convert prompt to embedding
 
198
  print("Retrieved Text ", retrieved_texts)
199
 
200
  return retrieved_texts
201
+ @spaces.GPU
202
  def respond(
203
  query,
204
  history: list[tuple[str, str]],