v

README.md

@@ -1,8 +1,8 @@
 ---
 title: Incoder Api
 emoji: 💻
-colorFrom: red
-colorTo: pink
+colorFrom: blue
+colorTo: green
 sdk: gradio
 sdk_version: 2.9.4
 app_file: app.py

app.py

@@ -2,18 +2,18 @@ import requests
 import os
 import gradio as gr
 import json
+from transformers import AutoModelForCausalLM, AutoTokenizer
 
-def start_server():   
-    os.system("uvicorn server:app --port 8080 --host 0.0.0.0 --workers 1")
-    os.environ["SSTART"] = "1"
-    
-if os.environ.get('SSTART') != "1":
-    start_server()     
+model_name = 'facebook/incoder-1B'
+tokenizer = AutoTokenizer.from_pretrained(model_name)
+model = AutoModelForCausalLM.from_pretrained(model_name, low_cpu_mem_usage=True)
+print('load ok')
 
-def completion(prompt,max_tokens,temperature,top_k,top_p):
-    req = f"http://0.0.0.0:8080?input_text={prompt}&top_p={top_p}&top_k={top_k}&temperature={temperature}&max_length={prompt}"
-    g = requests.get(req).json()
-    return g['text']
+def completion(prompt, max_tokens, temperature, top_k, top_p):
+    inpt = tokenizer.encode(prompt, return_tensors="pt")
+    out = model.generate(inpt, max_length=max_tokens, top_p=top_p, top_k=top_k, temperature=temperature)
+    res = tokenizer.decode(out[0])
+    return res
 
 demo = gr.Interface(
        fn=completion, 
@@ -32,5 +32,4 @@ demo = gr.Interface(
  
 )
 
-if __name__ == "__main__":
-    demo.launch()
+demo.launch()