prompt

app.py

@@ -1,5 +1,5 @@
 import spaces
-from transformers import pipeline, BitsAndBytesConfig, AutoModelForCausalLM, AutoTokenizer
+from transformers import pipeline
 import gradio as gr
 import torch
 import logging
@@ -12,24 +12,23 @@ try:
 except Exception as e:
     logging.warning(f"Could not import bitsandbytes: {e}")
 
-bnb_config = BitsAndBytesConfig(
-    load_in_4bit=True,
-    bnb_4bit_compute_dtype=torch.float16,  # key
-)
 
-tokenizer = AutoTokenizer.from_pretrained("ejschwartz/decaf-v1-22b-4bit")
-model = AutoModelForCausalLM.from_pretrained(
-    "ejschwartz/decaf-v1-22b-4bit",
-    device_map="auto",
-    quantization_config=bnb_config,
-)
+def decomp_create_prompt(input_data: str) -> str:
+    before = "# This is the decompiled code:\n"
+    after = "\n# What is the source code?\n"
+    prompt = before + input_data.strip() + after
 
-pipe = pipeline("text-generation", model=model, tokenizer=tokenizer)
+    return prompt
+
+
+pipe = pipeline(model="ejschwartz/decaf-v1-22b-4bit")
+pipe.model.to("cuda")
 
 @spaces.GPU(size="xlarge")
 def generate(text):
     print(f"Generating text... {text}")
-    return pipe(text)[0]
+    prompt = decomp_create_prompt(text)
+    return pipe(prompt)[0]['generated_text']
 
 demo = gr.Interface(fn=generate, inputs="text", outputs="text")
 demo.launch()