nyasukun commited on
Commit
fc804b6
·
1 Parent(s): fd96719
Files changed (1) hide show
  1. app.py +17 -20
app.py CHANGED
@@ -16,42 +16,39 @@ model_options = {
16
  "Foundation-Sec-8B": pipeline("text-generation", model="fdtn-ai/Foundation-Sec-8B"),
17
  }
18
 
19
- # Define the response function
20
- @spaces.GPU
21
  def generate_text_local(model_pipeline, prompt):
22
  """Local text generation"""
23
  try:
24
- logger.info(f"Running local text generation with {model_pipeline.path}")
25
-
 
 
26
  # Move model to GPU (entire pipeline)
27
  device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
28
- model_pipeline.model = model_pipeline.model.to(device)
29
-
30
- # Set other pipeline components to use GPU
31
- if hasattr(model_pipeline, "device"):
32
- model_pipeline.device = device
33
-
34
  # Record device information
35
- device_info = next(model_pipeline.model.parameters()).device
36
- logger.info(f"Model {model_pipeline.path} is running on device: {device_info}")
37
-
38
  outputs = model_pipeline(
39
  prompt,
40
- max_new_tokens=3, # = model.generate(max_new_tokens=3, …)
41
  do_sample=True,
42
  temperature=0.1,
43
  top_p=0.9,
44
- clean_up_tokenization_spaces=True, # echo 部分を整形
45
  )
46
 
47
  # Move model back to CPU
48
- model_pipeline.model = model_pipeline.model.to("cpu")
49
- if hasattr(model_pipeline, "device"):
50
- model_pipeline.device = torch.device("cpu")
51
-
52
  return outputs[0]["generated_text"].replace(prompt, "").strip()
53
  except Exception as e:
54
- logger.error(f"Error in local text generation with {model_pipeline.path}: {str(e)}")
55
  return f"Error: {str(e)}"
56
 
57
  # Build Gradio app
 
16
  "Foundation-Sec-8B": pipeline("text-generation", model="fdtn-ai/Foundation-Sec-8B"),
17
  }
18
 
19
+ #@spaces.GPU
 
20
  def generate_text_local(model_pipeline, prompt):
21
  """Local text generation"""
22
  try:
23
+ # モデル名取得(なければ 'unknown')
24
+ model_name = getattr(getattr(model_pipeline, "model", None), "name_or_path", "unknown")
25
+ logger.info(f"Running local text generation with {model_name}")
26
+
27
  # Move model to GPU (entire pipeline)
28
  device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
29
+ if hasattr(model_pipeline, "model"):
30
+ model_pipeline.model = model_pipeline.model.to(device)
31
+
 
 
 
32
  # Record device information
33
+ device_info = next(model_pipeline.model.parameters()).device if hasattr(model_pipeline, "model") else "unknown"
34
+ logger.info(f"Model {model_name} is running on device: {device_info}")
35
+
36
  outputs = model_pipeline(
37
  prompt,
38
+ max_new_tokens=3,
39
  do_sample=True,
40
  temperature=0.1,
41
  top_p=0.9,
42
+ clean_up_tokenization_spaces=True,
43
  )
44
 
45
  # Move model back to CPU
46
+ if hasattr(model_pipeline, "model"):
47
+ model_pipeline.model = model_pipeline.model.to("cpu")
48
+
 
49
  return outputs[0]["generated_text"].replace(prompt, "").strip()
50
  except Exception as e:
51
+ logger.error(f"Error in local text generation with {model_name}: {str(e)}")
52
  return f"Error: {str(e)}"
53
 
54
  # Build Gradio app