Update app.py

app.py

@@ -14,6 +14,8 @@ from chromadb.utils import embedding_functions
 from FlagEmbedding import BGEM3FlagModel
 from sklearn.metrics.pairwise import cosine_similarity
 
+device = "cuda" if torch.cuda.is_available() else "cpu"
+
 model = BGEM3FlagModel('BAAI/bge-m3',  
                        use_fp16=True) # Setting use_fp16 to True speeds up computation with a slight performance degradation
 
@@ -22,16 +24,16 @@ embeddings_data = pd.read_json("embeddings_tchap.json")
 embeddings_text = embeddings_data["text_with_context"].tolist()
 
 # Define the device
-#device = "cuda" if torch.cuda.is_available() else "cpu"
-#Define variables 
 temperature=0.2
 max_new_tokens=1000
 top_p=0.92
 repetition_penalty=1.7
 
-#model_name = "Pclanglais/Tchap"
+model_name = "Pclanglais/Tchap"
 
-#llm = LLM(model_name, max_model_len=4096)
+tokenizer = AutoTokenizer.from_pretrained(model_name)
+model = AutoModelForCausalLM.from_pretrained(model_name, torch_dtype=torch.float16)
+model = model.to('cuda:0')
 
 system_prompt = "<|begin_of_text|><|start_header_id|>system<|end_header_id|>\n\nTu es Albert, l'agent conversationnel des services publics qui peut décrire des documents de référence ou aider à des tâches de rédaction<|eot_id|>"
 
@@ -78,7 +80,7 @@ def predict(message, history):
 
     messages = system_prompt + messages
 
-    """"model_inputs = tokenizer([messages], return_tensors="pt").to("cuda")
+    model_inputs = tokenizer([messages], return_tensors="pt").to("cuda")
     streamer = TextIteratorStreamer(tokenizer, timeout=10., skip_prompt=True, skip_special_tokens=True)
     generate_kwargs = dict(
         model_inputs,
@@ -98,7 +100,7 @@ def predict(message, history):
     for new_token in streamer:
         if new_token != '<':
             partial_message += new_token
-            yield partial_message"""
+            yield partial_message
     return messages
 
 # Define the Gradio interface