adicionado chat

main.py

@@ -5,7 +5,7 @@ from typing import List, Literal
 from fastapi import FastAPI
 from pydantic import BaseModel
 from enum import Enum
-from transformers import M2M100Tokenizer, M2M100ForConditionalGeneration
+from transformers import M2M100Tokenizer, M2M100ForConditionalGeneration, AutoTokenizer, AutoModelForSeq2SeqLM
 import torch
 
 app = FastAPI(docs_url="/")
@@ -57,6 +57,45 @@ async def translate(request: TranslationRequest):
     return response
 
 
+
+# chat
+WHITESPACE_HANDLER = lambda k: re.sub('\s+', ' ', re.sub('\n+', ' ', k.strip()))
+ 
+chat_model_name = "csebuetnlp/mT5_multilingual_XLSum"
+tokenizer = AutoTokenizer.from_pretrained(chat_model_name)
+modelchat = AutoModelForSeq2SeqLM.from_pretrained(chat_model_name)
+
+@app.get("/chat")
+async def read_root(text: str, ):
+    input_ids = tokenizer(
+        [WHITESPACE_HANDLER(text)],
+        return_tensors="pt",
+        padding="max_length",
+        truncation=True,
+        max_length=512
+    )["input_ids"]
+
+        # max_length=84,
+    output_ids = modelchat.generate(
+        input_ids=input_ids,
+        max_length=500,
+        no_repeat_ngram_size=2,
+        num_beams=4
+    )[0]
+
+    summary = tokenizer.decode(
+        output_ids,
+        skip_special_tokens=True,
+        clean_up_tokenization_spaces=True
+    )
+
+    return {"summary": summary}
+
+
+
+
+
+
 if __name__ == "__main__":
     import uvicorn
 

requirements.txt

@@ -4,4 +4,11 @@ requests==2.27.*
 sentencepiece==0.1.*
 torch==1.11.*
 transformers==4.*
-uvicorn[standard]==0.17.*
+uvicorn[standard]==0.17.*
+ 
+tensorboard
+scikit-learn
+seqeval
+psutil
+sacrebleu 
+protobuf