added project files

Dockerfile

@@ -0,0 +1,10 @@
+FROM python:3.10-slim
+
+WORKDIR /app
+
+COPY requirements.txt .
+RUN pip install --no-cache-dir -r requirements.txt
+
+COPY . .
+
+CMD ["uvicorn", "app.main:app", "--host", "0.0.0.0", "--port", "7860"]

app/main.py

@@ -0,0 +1,45 @@
+from fastapi import FastAPI, Body
+from pydantic import BaseModel
+from app import models, prompts
+
+app = FastAPI()
+
+class TextInput(BaseModel):
+    text: str
+    mode: str
+    tone: str = None
+    target_lang: str = None
+
+@app.post("/rewrite")
+def rewrite(input: TextInput):
+    text = input.text
+    mode = input.mode
+
+    if mode == "grammar":
+        return {"result": models.run_grammar_correction(text)}
+
+    elif mode == "paraphrase":
+        prompt = prompts.paraphrase_prompt(text)
+        return {"result": models.run_flan_prompt(prompt)}
+
+    elif mode == "clarity":
+        prompt = prompts.clarity_prompt(text)
+        return {"result": models.run_flan_prompt(prompt)}
+
+    elif mode == "fluency":
+        prompt = prompts.fluency_prompt(text)
+        return {"result": models.run_flan_prompt(prompt)}
+
+    elif mode == "tone" and input.tone:
+        prompt = prompts.tone_prompt(text, input.tone)
+        return {"result": models.run_flan_prompt(prompt)}
+
+    elif mode == "translate" and input.target_lang:
+        return {"result": models.run_translation(text, input.target_lang)}
+
+    elif mode == "pronoun":
+        prompt = prompts.pronoun_friendly_prompt(text)
+        return {"result": models.run_flan_prompt(prompt)}
+
+    else:
+        return {"error": "Invalid request"}

app/models.py

@@ -0,0 +1,31 @@
+from transformers import AutoModelForSeq2SeqLM, AutoTokenizer
+import torch
+
+device = torch.device("cpu")
+
+# Grammar model
+grammar_tokenizer = AutoTokenizer.from_pretrained("vennify/t5-base-grammar-correction")
+grammar_model = AutoModelForSeq2SeqLM.from_pretrained("vennify/t5-base-grammar-correction").to(device)
+
+# FLAN-T5 for all prompts
+flan_tokenizer = AutoTokenizer.from_pretrained("google/flan-t5-small")
+flan_model = AutoModelForSeq2SeqLM.from_pretrained("google/flan-t5-small").to(device)
+
+# Translation model
+trans_tokenizer = AutoTokenizer.from_pretrained("Helsinki-NLP/opus-mt-en-ROMANCE")
+trans_model = AutoModelForSeq2SeqLM.from_pretrained("Helsinki-NLP/opus-mt-en-ROMANCE").to(device)
+
+def run_grammar_correction(text: str):
+    inputs = grammar_tokenizer(f"fix: {text}", return_tensors="pt").to(device)
+    outputs = grammar_model.generate(**inputs)
+    return grammar_tokenizer.decode(outputs[0], skip_special_tokens=True)
+
+def run_flan_prompt(prompt: str):
+    inputs = flan_tokenizer(prompt, return_tensors="pt").to(device)
+    outputs = flan_model.generate(**inputs)
+    return flan_tokenizer.decode(outputs[0], skip_special_tokens=True)
+
+def run_translation(text: str, target_lang: str):
+    inputs = trans_tokenizer(f">>{target_lang}<< {text}", return_tensors="pt").to(device)
+    outputs = trans_model.generate(**inputs)
+    return trans_tokenizer.decode(outputs[0], skip_special_tokens=True)

app/prompts.py

@@ -0,0 +1,14 @@
+def tone_prompt(text, tone):
+    return f"Rewrite the following text in a {tone} tone: {text}"
+
+def clarity_prompt(text):
+    return f"Make this clearer: {text}"
+
+def fluency_prompt(text):
+    return f"Improve the fluency of this sentence: {text}"
+
+def paraphrase_prompt(text):
+    return f"Paraphrase: {text}"
+
+def pronoun_friendly_prompt(text):
+    return f"Rewrite the text using inclusive and non-offensive pronouns: {text}"

requirements.txt

@@ -0,0 +1,8 @@
+fastapi
+uvicorn[standard]
+transformers
+torch
+sentencepiece
+pyspellchecker
+spacy
+nltk