ejschwartz commited on
Commit
83b6b2f
·
1 Parent(s): 87253bd
Files changed (1) hide show
  1. app.py +10 -6
app.py CHANGED
@@ -12,9 +12,7 @@ print("Downloading model")
12
 
13
  tokenizer = AutoTokenizer.from_pretrained(model_id, gguf_file=filename)
14
  model = AutoModelForCausalLM.from_pretrained(
15
- model_id,
16
- gguf_file=filename,
17
- device_map="auto"
18
  )
19
 
20
  # Then create the pipeline with the model and tokenizer
@@ -33,6 +31,7 @@ system = """<s>[INST]You are an expert at analyzing code that has been decompile
33
  Analyze the following IDA Hex Rays pseudocode and generate a valid JSON object containing the keys 'function_name','comment', and an array 'variables' explaining what the code does, suggest a function name based on the analysis of the code, and new variable names based on the analysis of the code.[/INST]</s>
34
  """
35
 
 
36
  @spaces.GPU
37
  def predict(code):
38
  prompt = f"""{system}
@@ -45,6 +44,11 @@ def predict(code):
45
  print(f"Tokenized: {tokenizer.tokenize(prompt)}")
46
  return pipe(prompt)
47
 
48
- demo = gr.Interface(fn=predict, inputs="text", outputs="text",
49
- description=frontmatter.load("README.md").content)
50
- demo.launch()
 
 
 
 
 
 
12
 
13
  tokenizer = AutoTokenizer.from_pretrained(model_id, gguf_file=filename)
14
  model = AutoModelForCausalLM.from_pretrained(
15
+ model_id, gguf_file=filename, device_map="auto"
 
 
16
  )
17
 
18
  # Then create the pipeline with the model and tokenizer
 
31
  Analyze the following IDA Hex Rays pseudocode and generate a valid JSON object containing the keys 'function_name','comment', and an array 'variables' explaining what the code does, suggest a function name based on the analysis of the code, and new variable names based on the analysis of the code.[/INST]</s>
32
  """
33
 
34
+
35
  @spaces.GPU
36
  def predict(code):
37
  prompt = f"""{system}
 
44
  print(f"Tokenized: {tokenizer.tokenize(prompt)}")
45
  return pipe(prompt)
46
 
47
+
48
+ demo = gr.Interface(
49
+ fn=predict,
50
+ inputs="text",
51
+ outputs="text",
52
+ description=frontmatter.load("README.md").content,
53
+ )
54
+ demo.launch()