Update app.py

app.py

@@ -9,11 +9,28 @@ from text_generation import Client
 
 from share_btn import community_icon_html, loading_icon_html, share_js, share_btn_css
 
-HF_TOKEN = os.environ.get("HF_TOKEN", None)
+#HF_TOKEN = os.environ.get("HF_TOKEN", None)
+
+#API_URL = "https://api-inference.huggingface.co/models/bigcode/starcoder"
+#API_URL_BASE ="https://api-inference.huggingface.co/models/bigcode/starcoderbase"
+#API_URL_PLUS = "https://api-inference.huggingface.co/models/bigcode/starcoderplus"
+https://huggingface.co/smallcloudai/Refact-1_6B-fim/discussions
+
+from transformers import AutoModelForCausalLM, AutoTokenizer
+
+checkpoint = "smallcloudai/Refact-1_6B-fim"
+device = "cuda" # for GPU usage or "cpu" for CPU usage
+
+tokenizer = AutoTokenizer.from_pretrained(checkpoint)
+model = AutoModelForCausalLM.from_pretrained(checkpoint, trust_remote_code=True).to(device)
+
+prompt = '<fim_prefix>def print_hello_world():\n    """<fim_suffix>\n    print("Hello world!")<fim_middle>'
+
+inputs = tokenizer.encode(prompt, return_tensors="pt").to(device)
+outputs = model.generate(inputs, max_length=100, temperature=0.2)
+print("-"*80)
+print(tokenizer.decode(outputs[0]))
 
-API_URL = "https://api-inference.huggingface.co/models/bigcode/starcoder"
-API_URL_BASE ="https://api-inference.huggingface.co/models/bigcode/starcoderbase"
-API_URL_PLUS = "https://api-inference.huggingface.co/models/bigcode/starcoderplus"
 
 FIM_PREFIX = "<fim_prefix>"
 FIM_MIDDLE = "<fim_middle>"
@@ -75,16 +92,8 @@ theme = gr.themes.Monochrome(
     ],
 )
 
-client = Client(
-    API_URL,
-    headers={"Authorization": f"Bearer {HF_TOKEN}"},
-)
-client_base = Client(
-    API_URL_BASE, headers={"Authorization": f"Bearer {HF_TOKEN}"},
-)
-client_plus = Client(
-    API_URL_PLUS, headers={"Authorization": f"Bearer {HF_TOKEN}"},
-)
+inputs = tokenizer.encode(prompt, return_tensors="pt").to(device)
+outputs = model.generate(inputs, max_length=100, temperature=0.2)
 
 def generate(
     prompt, temperature=0.9, max_new_tokens=256, top_p=0.95, repetition_penalty=1.0, version="StarCoder",
@@ -113,29 +122,9 @@ def generate(
             raise ValueError(f"Only one {FIM_INDICATOR} allowed in prompt!")
         prompt = f"{FIM_PREFIX}{prefix}{FIM_SUFFIX}{suffix}{FIM_MIDDLE}"
 
-    if version == "StarCoder":
-        stream = client.generate_stream(prompt, **generate_kwargs)
-    elif version == "StarCoderPlus":
-        stream = client_plus.generate_stream(prompt, **generate_kwargs)
-    else:
-        stream = client_base.generate_stream(prompt, **generate_kwargs)
-
-    if fim_mode:
-        output = prefix
-    else:
-        output = prompt
-
-    previous_token = ""
-    for response in stream:
-        if response.token.text == "<|endoftext|>":
-            if fim_mode:
-                output += suffix
-            else:
-                return output
-        else:
-            output += response.token.text
-        previous_token = response.token.text
-        yield output
+    inputs = tokenizer.encode(prompt, return_tensors="pt").to(device)
+    output = model.generate(inputs, max_length=100, temperature=0.2)
+
     return output