use different app for testing

Dockerfile

@@ -6,12 +6,9 @@ COPY . .
 
 # Set the working directory to /
 WORKDIR /
-VOLUME ["cache:/.cache"]
 
 # Install requirements.txt 
 RUN pip install --no-cache-dir --upgrade -r /requirements.txt
 
-USER 1000
-
 # Start the FastAPI app on port 7860, the default port expected by Spaces
 CMD ["uvicorn", "app:app", "--host", "0.0.0.0", "--port", "7860"]

app.py

@@ -1,17 +1,35 @@
-from fastapi import FastAPI
-from typing import List
-import torch
-from transformers import AutoModelForSeq2SeqLM, AutoTokenizer
-from IndicTransToolkit import IndicProcessor
+from pydantic import BaseModel
+
+from .ConfigEnv import config
 from fastapi.middleware.cors import CORSMiddleware
 
-import os
+from langchain.llms import Clarifai
+from langchain.chains import LLMChain
+from langchain.prompts import PromptTemplate
+from TextGen import app
+
+
+class Generate(BaseModel):
+    text: str
+
+
+def generate_text(prompt: str):
+    if prompt == "":
+        return {"detail": "Please provide a prompt."}
+    else:
+        prompt = PromptTemplate(template=prompt, input_variables=["Prompt"])
+        llm = Clarifai(
+            pat=config.CLARIFAI_PAT,
+            user_id=config.USER_ID,
+            app_id=config.APP_ID,
+            model_id=config.MODEL_ID,
+            model_version_id=config.MODEL_VERSION_ID,
+        )
+        llmchain = LLMChain(prompt=prompt, llm=llm)
+        llm_response = llmchain.run({"Prompt": prompt})
+        return Generate(text=llm_response)
 
-os.environ["HF_HOME"] = "/.cache"
-# Initialize FastAPI
-app = FastAPI()
 
-# Add CORS middleware
 app.add_middleware(
     CORSMiddleware,
     allow_origins=["*"],
@@ -20,67 +38,17 @@ app.add_middleware(
     allow_headers=["*"],
 )
 
-# Initialize models and processors
-model = AutoModelForSeq2SeqLM.from_pretrained(
-    "ai4bharat/indictrans2-en-indic-1B", trust_remote_code=True
-)
-tokenizer = AutoTokenizer.from_pretrained(
-    "ai4bharat/indictrans2-en-indic-1B", trust_remote_code=True
-)
-ip = IndicProcessor(inference=True)
-DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
-model = model.to(DEVICE)
-
-
-def translate_text(sentences: List[str], target_lang: str):
-    try:
-        src_lang = "eng_Latn"
-        batch = ip.preprocess_batch(sentences, src_lang=src_lang, tgt_lang=target_lang)
-        inputs = tokenizer(
-            batch,
-            truncation=True,
-            padding="longest",
-            return_tensors="pt",
-            return_attention_mask=True,
-        ).to(DEVICE)
-
-        with torch.no_grad():
-            generated_tokens = model.generate(
-                **inputs,
-                use_cache=True,
-                min_length=0,
-                max_length=256,
-                num_beams=5,
-                num_return_sequences=1,
-            )
-
-        with tokenizer.as_target_tokenizer():
-            generated_tokens = tokenizer.batch_decode(
-                generated_tokens.detach().cpu().tolist(),
-                skip_special_tokens=True,
-                clean_up_tokenization_spaces=True,
-            )
 
-        translations = ip.postprocess_batch(generated_tokens, lang=target_lang)
-        return {
-            "translations": translations,
-            "source_language": src_lang,
-            "target_language": target_lang,
-        }
-    except Exception as e:
-        raise Exception(f"Translation failed: {str(e)}")
+@app.get("/", tags=["Home"])
+def api_home():
+    return {"detail": "Welcome to FastAPI TextGen Tutorial!"}
 
 
-# FastAPI routes
-@app.get("/health")
-async def health_check():
-    return {"status": "healthy"}
-
-
-@app.post("/translate")
-async def translate_endpoint(sentences: List[str], target_lang: str):
-    try:
-        result = translate_text(sentences=sentences, target_lang=target_lang)
-        return result
-    except Exception as e:
-        raise HTTPException(status_code=500, detail=str(e))
+@app.post(
+    "/api/generate",
+    summary="Generate text from prompt",
+    tags=["Generate"],
+    response_model=Generate,
+)
+def inference(input_prompt: str):
+    return generate_text(prompt=input_prompt)

requirements.txt

@@ -1,6 +1,7 @@
-fastapi
+fastapi==0.99.1
 uvicorn
-torch
-transformers
-git+https://github.com/VarunGumma/IndicTransToolkit.git
-python-multipart
+requests
+pydantic==1.10.12
+langchain
+clarifai
+Pillow