pratham0203 commited on
Commit
30f98e0
·
verified ·
1 Parent(s): f4233f3
Files changed (1) hide show
  1. app.py +6 -6
app.py CHANGED
@@ -5,23 +5,23 @@ from transformers import pipeline
5
  app=FastAPI()
6
 
7
  # Initialize the text generation pipeline
8
- pipe = pipeline("text-generation", model="meta-llama/Meta-Llama-3.1-8B")
9
 
10
 
11
  @app.get("/")
12
  def home():
13
- return {"message":"Welcome DFE team"}
14
 
15
  # Define a function to handle the GET request at `/generate`
16
 
17
 
18
  @app.get("/generate")
19
  def generate(text:str):
20
- # messages = [
21
- # {"role": "user", "content": text},
22
- # ]
23
  ## use the pipeline to generate text from given input text
24
- output=pipe(text)
25
 
26
  ## return the generate text in Json reposne
27
  return {"output":output[0]['generated_text']}
 
5
  app=FastAPI()
6
 
7
  # Initialize the text generation pipeline
8
+ pipe = pipeline("text-generation", model="deepseek-ai/DeepSeek-Coder-V2-Lite-Base", trust_remote_code=True)
9
 
10
 
11
  @app.get("/")
12
  def home():
13
+ return {"message":"Welcome DFE team"}
14
 
15
  # Define a function to handle the GET request at `/generate`
16
 
17
 
18
  @app.get("/generate")
19
  def generate(text:str):
20
+ messages = [
21
+ {"role": "user", "content": text},
22
+ ]
23
  ## use the pipeline to generate text from given input text
24
+ output=pipe(messages)
25
 
26
  ## return the generate text in Json reposne
27
  return {"output":output[0]['generated_text']}