move to one device

app.py

@@ -45,7 +45,7 @@ last_messages = []
 def load_model():
     ankerbot_model = AutoModelForCausalLM.from_pretrained(
         model_name,
-        device_map="cpu",
+        device_map="auto",
         torch_dtype=torch.float16,
         use_cache=True,
         offload_folder="../offload"
@@ -55,9 +55,8 @@ def load_model():
         torch_dtype=torch.float16,
         truncation=True,
         padding=True, )
-    ankerbot_model.to("cuda") 
-    generator = pipeline(task="text-generation", model=ankerbot_model, tokenizer=ankerbot_tokenizer, torch_dtype=torch.float16, trust_remote_code=False) # True for flash-attn2 else False
-    generator_mini = pipeline(task="text-generation", model=ankerbot_model, tokenizer=ankerbot_tokenizer, torch_dtype=torch.float16, trust_remote_code=False) # True for flash-attn2 else False
+    generator = pipeline(task="text-generation", model=ankerbot_model, tokenizer=ankerbot_tokenizer, torch_dtype=torch.float16, trust_remote_code=False, device="cuda:0") # True for flash-attn2 else False
+    generator_mini = pipeline(task="text-generation", model=ankerbot_model, tokenizer=ankerbot_tokenizer, torch_dtype=torch.float16, trust_remote_code=False, device="cuda:0") # True for flash-attn2 else False
     return (generator, generator_mini)
     
 _model_cache = None
@@ -81,9 +80,9 @@ def generate_response(query, context, prompts, max_tokens, temperature, top_p, g
         Du bekommst Kundenanfragen zum Beispiel zu einer Bestellung, antworte Anhand des zur Verfügunggestellten Kontextes.
         Tu so, als wär der Kontext Bestandteil deines Wissens. Sprich den Kunden persönlich an.
         Nenne nichts außerhalb des Kontext.
-        Konversation: {",".join(last_messages)}
         Kontext Kundenservice: {context}
         <|im_end|>
+        {"".join(last_messages)}
         <|im_start|>user
         Frage: {query}
         <|im_end|>
@@ -97,9 +96,9 @@ def generate_response(query, context, prompts, max_tokens, temperature, top_p, g
         Du bekommst im Kontext Informationen zu Rezepten und Gerichten.
         Tu so, als wär der Kontext Bestandteil deines Wissens. Sprich den Kunden persönlich an.
         Nenne nichts außerhalb des Kontext.
-        Konversation: {",".join(last_messages)}
         Kontext Rezepte: {context}
         <|im_end|>
+        {"".join(last_messages)}
         <|im_start|>user
         Frage: {query}
         <|im_end|>
@@ -113,9 +112,9 @@ def generate_response(query, context, prompts, max_tokens, temperature, top_p, g
         Du bekommst im Kontext Informationen zu Produkte, nach denen gefragt ist, oder welche ähnlich sein könnten.
         Tu so, als wär der Kontext Bestandteil deines Wissens. Sprich den Kunden persönlich an.
         Nenne nichts außerhalb des Kontext.
-        Konversation: {",".join(last_messages)}
         Kontext Produkte: {context}
         <|im_end|>
+        {"".join(last_messages)}
         <|im_start|>user
         Frage: {query}
         <|im_end|>
@@ -173,6 +172,7 @@ def respond(
             Frage: {query}
         <|im_end|>
         <|im_start|>assistant"""
+    system_message = system_message.to("cuda") if torch.cuda.is_available() else system_message
     refined_context = generator[1](system_message, do_sample=True, padding=True, truncation=True, top_p=0.95, max_new_tokens=150)
     # Retrieve relevant context from Qdrant
     if "rezept" in query.lower() or "gericht" in query.lower():
@@ -184,7 +184,7 @@ def respond(
 
     context = search_qdrant_with_context(query + " " + refined_context[0]["generated_text"].split("assistant\n").pop(), collection_name)
     answer = generate_response(query, context, last_messages, max_tokens, temperature, top_p, generator[0])
-    full_conv = f"Nutzer:{query};Assistent:{answer}"
+    full_conv = f"<|im_start|>user {query}<|im_end|><|im_start|>assistent {answer}<|im_end|>"
     if len(last_messages) > 5:
         last_messages.pop(0)
     last_messages.append(full_conv)