v

app.py

@@ -1,16 +1,11 @@
 import streamlit as st
-from transformers import AutoModelForCausalLM, AutoTokenizer
-
-@st.cache(allow_output_mutation=True)
-def load_model():
-    '''
-    Loads the model and tokenizer from the local directory.
-    :return: A list containing the model and the tokenizer.
-    '''
-    model_name = 'facebook/incoder-1B'
-    tokenizer = AutoTokenizer.from_pretrained(model_name)
-    model = AutoModelForCausalLM.from_pretrained(model_name, low_cpu_mem_usage=True)
-    return [model, tokenizer]
+import requests
+import os
+import json
+
+def start_server():   
+    os.system("uvicorn server:app --port 8080 --host 0.0.0.0 --workers 2")
+    st.session_state['server_started'] = True
 
 st.set_page_config(
     page_title="Incoder Example",
@@ -37,11 +32,7 @@ repetition_penalty = st.slider('repetition_penalty', 1.0, 10.0, 1.0)
 
 if gen:
     c.code('Generating...')
-    m = load_model()
-
-    inpt = m[1].encode(txt, return_tensors="pt")
-    out = m[0].generate(inpt, max_length=max_length, top_p=top_p, top_k=top_k, temperature=temperature, num_beams=num_beams, repetition_penalty=repetition_penalty)
-    res = m[1].decode(out[0])
-
+    req = f"http://0.0.0.0:8080?input_text={txt}&top_p={top_p}&top_k={top_k}&temperature={temperature}&num_beams={num_beams}&repetition_penalty={repetition_penalty}&max_length={max_length}"
+    res = requests.get(req)
     print('ok')
-    c.code(res)
+    c.code(json.loads(res.text))

requirements.txt

@@ -1,3 +1,5 @@
 transformers
-streamlit
-torch
+fastapi
+uvicorn
+torch
+requests

server.py

@@ -0,0 +1,16 @@
+import uvicorn
+from fastapi import FastAPI
+from transformers import AutoModelForCausalLM, AutoTokenizer
+
+app = FastAPI()
+
+model_name = 'facebook/incoder-1B'
+tokenizer = AutoTokenizer.from_pretrained(model_name)
+model = AutoModelForCausalLM.from_pretrained(model_name, low_cpu_mem_usage=True)
+
+@app.get("/")
+def read_root(input_text, max_length, top_p, top_k, num_beams, temperature, repetition_penalty):
+    inpt = model_name.encode(input_text, return_tensors="pt")
+    out = model.generate(inpt, max_length=max_length, top_p=top_p, top_k=top_k, temperature=temperature, num_beams=num_beams, repetition_penalty=repetition_penalty)
+    res = tokenizer.decode(out[0])
+    return {res}