Update app.py

app.py

@@ -217,11 +217,11 @@ def simple_chat(message: dict, temperature: float = 0.8, max_length: int = 4096,
     try:
         model = AutoModelForCausalLM.from_pretrained(
             MODEL_ID,
-            torch_dtype=torch.bfloat16,
+            torch_dtype=torch.bfloat16 if torch.cuda.is_available() else torch.float32,
             low_cpu_mem_usage=True,
             trust_remote_code=True
         )
-        
+
         tokenizer = AutoTokenizer.from_pretrained(MODEL_ID, trust_remote_code=True)
 
         conversation = []
@@ -230,14 +230,14 @@ def simple_chat(message: dict, temperature: float = 0.8, max_length: int = 4096,
         if "file_content" in message and message["file_content"]:
             file_content = message["file_content"]
             file_name = message["file_name"]
-            
+
             # Guardar el archivo en un archivo temporal
             with open(file_name, "wb") as f:
                 f.write(file_content.read())
-            
+
             # Llamar a `mode_load` con el nombre del archivo
             choice, contents = mode_load(file_name)
-            
+
             if choice == "image":
                 conversation.append({"role": "user", "image": contents, "content": message['text']})
             elif choice == "doc":
@@ -267,31 +267,31 @@ def simple_chat(message: dict, temperature: float = 0.8, max_length: int = 4096,
             eos_token_id=[151329, 151336, 151338],
         )
 
-        # gen_kwargs = {**input_ids, **generate_kwargs}
+        gen_kwargs = {**input_ids, **generate_kwargs}
 
-        # with torch.no_grad():
-        #     thread = Thread(target=model.generate, kwargs=gen_kwargs)
-        #     thread.start()
-        #     buffer = ""
-        #     for new_text in streamer:
-        #         buffer += new_text
-        #         yield buffer
+        # Define the function to run generation
+        def generate_text():
+            with torch.no_grad():
+                model.generate(**gen_kwargs, streamer=streamer)
 
-        # print("--------------")
-        # print("Buffer: ")
-        # print("  ")
-        # print(buffer)
-        # print("  ")
-        # print("--------------")
-
-        with torch.no_grad():
-            generated_ids = model.generate(input_ids['input_ids'], **generate_kwargs)
-            generated_text = tokenizer.decode(generated_ids[0], skip_special_tokens=True)
+        # Start the generation in a separate thread
+        thread = Thread(target=generate_text)
+        thread.start()
 
-        #generated_text = buffer
+        def stream_response():
+            buffer = ""
+            for new_text in streamer:
+                buffer += new_text
+                yield new_text
+            print("--------------")
+            print("Buffer: ")
+            print("  ")
+            print(buffer)
+            print("  ")
+            print("--------------")
 
+        return StreamingResponse(stream_response(), media_type="text/plain")
 
-        return PlainTextResponse(generated_text)
     except Exception as e:
         return PlainTextResponse(f"Error: {str(e)}")