HTAR5 commited on
Commit
4258af5
1 Parent(s): 248e15e

Rename appmod.py to app.py

Browse files
Files changed (1) hide show
  1. appmod.py → app.py +30 -30
appmod.py → app.py RENAMED
@@ -1,30 +1,30 @@
1
- from fastapi import FastAPI, Query
2
- from transformers import pipeline
3
-
4
- app = FastAPI()
5
-
6
- def initialize_pipeline():
7
- return pipeline("text2text-generation", model="google/flan-t5-small")
8
-
9
- pipe = initialize_pipeline()
10
-
11
- @app.get("/")
12
- def home():
13
- return {"message": "Hello Siddhant"}
14
-
15
- @app.get("/generate")
16
- def generate_text(
17
- text: str = Query(None, description="Input text to generate from"),
18
- prompt: str = Query(None, description="Optional prompt for fine-tuning the generated text"),
19
- ):
20
- if not text and not prompt:
21
- return {"error": "Please provide either 'text' or 'prompt' parameter."}
22
-
23
- if prompt:
24
- input_text = f"{text} {prompt}" if text else prompt
25
- else:
26
- input_text = text
27
-
28
- output = pipe(input_text, max_length=100, do_sample=True, top_k=50)
29
-
30
- return {"input_text": input_text, "output": output[0]["generated_text"]}
 
1
+ from fastapi import FastAPI, Query
2
+ from transformers import pipeline
3
+
4
+ app = FastAPI()
5
+
6
+ def initialize_pipeline():
7
+ return pipeline("text-classification", model="FacebookAI/roberta-large-mnli")
8
+
9
+ pipe = initialize_pipeline()
10
+
11
+ # @app.get("/docs")
12
+ # def home():
13
+ # return {"message": "Hello Siddhant"}
14
+
15
+ @app.get("/")
16
+ def generate_text(
17
+ text: str = Query(None, description="Input text to generate from"),
18
+ prompt: str = Query(None, description="Optional prompt for fine-tuning the generated text"),
19
+ ):
20
+ if not text and not prompt:
21
+ return {"error": "Please provide either 'text' or 'prompt' parameter."}
22
+
23
+ if prompt:
24
+ input_text = f"{text} {prompt}" if text else prompt
25
+ else:
26
+ input_text = text
27
+
28
+ output = pipe(input_text, max_length=100, do_sample=True, top_k=50)
29
+
30
+ return {"input_text": input_text, "output": output[0]["generated_text"]}