Update app.py

app.py

@@ -273,95 +273,18 @@ def simple_chat(message, temperature: float = 0.8, max_length: int = 4096, top_p
 
     thread.join()
 
-    print("---------")
-    print("Text: ")
-    print(buffer)
-    print("---------")
-
-        # If the generated text is JSON, parse it and format it as plain text
     try:
         json_content = json.loads(buffer)
         formatted_text = "\n".join(f"{key}: {value}" for key, value in json_content.items())
     except json.JSONDecodeError:
-        # If the buffer is not valid JSON, return it as is
         formatted_text = buffer
 
-    return PlainTextResponse(formatted_text)
-
-
-
-
-# @spaces.GPU()
-# def simple_chat(message: dict, temperature: float = 0.8, max_length: int = 4096, top_p: float = 1, top_k: int = 10, penalty: float = 1.0):
-#     try:
-#         model = AutoModelForCausalLM.from_pretrained(
-#             MODEL_ID,
-#             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 = []
-
-#         if "file_content" in message and message["file_content"]:
-#             file_content = message["file_content"]
-#             file_name = message["file_name"]
-
-#             with open(file_name, "wb") as f:
-#                 f.write(file_content.read())
-
-#             choice, contents = mode_load(file_name)
-
-#             if choice == "image":
-#                 conversation.append({"role": "user", "image": contents, "content": message['text']})
-#             elif choice == "doc":
-#                 message['text'] = contents + "\n\n\n" + "{} files uploaded.\n".format(1) + message['text']
-#                 conversation.append({"role": "user", "content": message['text']})
-#                 # format_msg = contents + "\n\n\n" + "{} files uploaded.\n".format(1) + message['text']
-#                 # conversation.append({"role": "user", "content": format_msg})
-#         else:
-#             conversation.append({"role": "user", "content": message['text']})
-
-#         input_ids = tokenizer.apply_chat_template(conversation, tokenize=True, add_generation_prompt=True, return_tensors="pt", return_dict=True).to(model.device)
-
-#         streamer = TextIteratorStreamer(tokenizer, timeout=60.0, skip_prompt=True, skip_special_tokens=True)
-
-#         generate_kwargs = dict(
-#             max_length=max_length,
-#             do_sample=True,
-#             top_p=top_p,
-#             top_k=top_k,
-#             temperature=temperature,
-#             repetition_penalty=penalty,
-#             eos_token_id=[151329, 151336, 151338],
-#         )
-
-#         gen_kwargs = {**input_ids, **generate_kwargs}
-
-#         for entry in conversation:
-#             print(f"Role: {entry['role']}, Content: {entry.get('content', '')}")
-
-#         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)
-
-#         text_original = message['text'].strip()
-#         generated_text_cleaned = generated_text.replace(text_original, "").strip()
-
-#         print(" ")
-#         print("---------")
-#         print("Text: ")
-#         print("  ")
-#         print(generated_text_cleaned)
-
-
-#         return PlainTextResponse(generated_text_cleaned)
-#     except Exception as e:
-#         return PlainTextResponse(f"Error: {str(e)}")
-
+    print("---------")
+    print("Text: ")
+    print(formatted_text)
+    print("---------")
 
+    return PlainTextResponse(formatted_text)
 
 
 @app.post("/chat/")