Rockramsri commited on
Commit
f17d901
1 Parent(s): 3035c06
Files changed (1) hide show
  1. app.py +7 -7
app.py CHANGED
@@ -1,13 +1,14 @@
1
  from fastapi import FastAPI
2
- from transformers import pipeline
3
- import os
4
 
5
  ## create a new FASTAPI app instance
6
  app=FastAPI()
7
 
8
  # Initialize the text generation pipeline
9
- pipe = pipeline("text2text-generation", model="lmstudio-community/Meta-Llama-3.1-8B-Instruct-GGUF",token=os.getenv('HF_KEY'))
10
-
 
 
11
 
12
  @app.get("/")
13
  def home():
@@ -19,9 +20,8 @@ def home():
19
  @app.get("/generate")
20
  def generate(text:str):
21
  ## use the pipeline to generate text from given input text
22
- output=pipe(text)
23
-
24
  ## return the generate text in Json reposne
25
- return {"output":output[0]['generated_text']}
26
 
27
 
 
1
  from fastapi import FastAPI
2
+ from llama_cpp import Llama
 
3
 
4
  ## create a new FASTAPI app instance
5
  app=FastAPI()
6
 
7
  # Initialize the text generation pipeline
8
+ #pipe = pipeline("text2text-generation", model="lmstudio-community/Meta-Llama-3.1-8B-Instruct-GGUF",token=os.getenv('HF_KEY'))
9
+ llm = Llama(
10
+ model_path="Meta-Llama-3.1-8B-Instruct-Q4_K_M.gguf",
11
+ )
12
 
13
  @app.get("/")
14
  def home():
 
20
  @app.get("/generate")
21
  def generate(text:str):
22
  ## use the pipeline to generate text from given input text
23
+ output== llm(text, max_tokens=1000)
 
24
  ## return the generate text in Json reposne
25
+ return {"output":output["choices"][0]["text"]}
26
 
27