Update app.py

app.py

@@ -1,7 +1,6 @@
 import os
 import gradio as gr
 import spaces
-import torch
 
 # 支持的语言选项
 LANGUAGES = {
@@ -100,41 +99,57 @@ LANGUAGES = {
     "Western Persian": "Western Persian"
 }
 
-from huggingface_hub import snapshot_download
-from vllm import LLM
-from vllm.sampling_params import BeamSearchParams
+from transformers import AutoModelForCausalLM, AutoTokenizer
 
-save_dir = "./model_weights"
-repo_id = "ByteDance-Seed/Seed-X-PPO-7B"
-cache_dir = save_dir + "/cache"
 
-snapshot_download(
-  cache_dir=cache_dir,
-  local_dir=save_dir,
-  repo_id=repo_id,
-  local_dir_use_symlinks=False,
-  resume_download=True,
-  allow_patterns=["*.json", "*.safetensors", "*.bin", "*.py", "*.md", "*.txt"],
-)
+# from huggingface_hub import snapshot_download
+# save_dir = "./model_weights"
+# repo_id = "ByteDance-Seed/Seed-X-PPO-7B"
+# cache_dir = save_dir + "/cache"
 
-device = torch.device("cuda")
-model = LLM(model=save_dir,
-                max_num_seqs=512,
-                # tensor_parallel_size=8,
-                enable_prefix_caching=True, 
-                gpu_memory_utilization=0.95).to(device)
+# snapshot_download(
+#   cache_dir=cache_dir,
+#   local_dir=save_dir,
+#   repo_id=repo_id,
+#   local_dir_use_symlinks=False,
+#   resume_download=True,
+#   allow_patterns=["*.json", "*.safetensors", "*.bin", "*.py", "*.md", "*.txt"],
+# )
+
+device = "cuda"
+MODEL_NAME = "ByteDance-Seed/Seed-X-PPO-7B"
+print("Start dowload")
+def load_model():
+    model = AutoModelForCausalLM.from_pretrained(MODEL_NAME,torch_dtype="bfloat16").to(device)
+    print(f"Model loaded in {device}")
+    return model
+
+model = load_model()
+print("Ednd dowload")
+# Loading the tokenizer once, because re-loading it takes about 1.5 seconds each time
+tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)
 
-decoding_params = BeamSearchParams(beam_width=4,
-                                   max_tokens=512)
 
 @spaces.GPU(duration=120)
 def translate_text(text, source_lang, target_lang):
     if not text.strip():
         return "请输入要翻译的文本"
     try:
-        response = model.generate(messages, decoding_params)
-        response = [res.outputs[0].text.strip() for res in results]
-        yield response[0]
+        input_tokens = (
+            tokenizer(text, return_tensors="pt")
+            .input_ids[0]
+            .cpu()
+            .numpy()
+            .tolist()
+        )
+        translated_chunk = model.generate(
+            input_ids=torch.tensor([input_tokens]).to(device),
+            max_length=len(input_tokens) + 2048,
+            num_return_sequences=1,
+        )
+        full_output = tokenizer.decode(translated_chunk[0], skip_special_tokens=True)
+        full_output = full_output.replace(input_text,"")
+        yield full_output
     except Exception as e:
         yield f"翻译出错: {str(e)}"